Update app.py

app.py

@@ -54,14 +54,15 @@ def load_model():
         app.logger.info(f"Device set to use {device}")
         
         model = AutoModelForCausalLM.from_pretrained(
-            "gpt2-medium",
+            "gpt2",
             use_safetensors=True,
             device_map="auto",
             torch_dtype=dtype,
-            low_cpu_mem_usage=True
+            low_cpu_mem_usage=True,
+            offload_folder="offload"
         )
         
-        tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
         
         # Initialize pipeline without explicit device assignment
         generator = pipeline(
@@ -146,17 +147,11 @@ IEEE_TEMPLATE = """
 # --------------------------------------------------
 @app.route('/health', methods=['GET'])
 def health_check():
-    app.logger.info("Health check requested")
-    
-    if load_error:
-        app.logger.error(f"Health check failed: {load_error}")
-        return jsonify({
-            "status": "error",
-            "message": f"Model failed to load: {load_error}"
-        }), 500
-        
-    status_code = 200 if model_loaded else 503
-    device_info = "cuda" if torch.cuda.is_available() else "cpu"
+    return jsonify({
+        "status": "ok",
+        "model_loaded": model_loaded,
+        "device": "cuda" if torch.cuda.is_available() else "cpu"
+    }), 200
     
     app.logger.info(f"Health check returning status: {'ready' if model_loaded else 'loading'}, device: {device_info}")
     return jsonify({