Update app.py

app.py

@@ -20,16 +20,18 @@ class CoachingInput(BaseModel):
 
 # Define model path (absolute path in the container)
 model_path = "/app/fine-tuned-construction-llm"
-
-# Verify the model directory exists
-if not os.path.isdir(model_path):
-    logger.error(f"Model directory not found: {model_path}")
-    raise Exception(f"Model directory not found: {model_path}")
+fallback_model = "gpt2"  # Fallback to a pre-trained model if local model is unavailable
 
 # Load model and tokenizer
 try:
-    model = AutoModelForCausalLM.from_pretrained(model_path, local_files_only=True)
-    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
+    if os.path.isdir(model_path):
+        logger.info(f"Loading local model from {model_path}")
+        model = AutoModelForCausalLM.from_pretrained(model_path, local_files_only=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
+    else:
+        logger.warning(f"Model directory not found: {model_path}. Falling back to pre-trained model: {fallback_model}")
+        model = AutoModelForCausalLM.from_pretrained(fallback_model)
+        tokenizer = AutoTokenizer.from_pretrained(fallback_model)
     logger.info("Model and tokenizer loaded successfully")
 except Exception as e:
     logger.error(f"Failed to load model or tokenizer: {str(e)}")
@@ -60,18 +62,24 @@ async def generate_coaching(data: CoachingInput):
         # Decode and parse response
         response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Simulate structured output (replace with actual parsing logic based on model output)
-        # This assumes the model outputs a JSON-like string; adjust based on fine-tuning
-        try:
-            response_json = json.loads(response_text)
-        except json.JSONDecodeError:
-            # Fallback: Construct a default response if parsing fails
-            response_json = {
-                "checklist": ["Inspect safety equipment", "Review milestone progress"],
-                "tips": ["Prioritize team communication", "Check weather updates"],
-                "quote": "Every step forward counts!"
-            }
-            logger.warning("Failed to parse model output as JSON, using default response")
+        # Since gpt2 may not output JSON, parse the response manually or use fallback
+        # This is a simplified parsing logic; adjust based on your model's output format
+        if not response_text.startswith("{"):
+            checklist = ["Inspect safety equipment", "Review milestone progress"]
+            tips = ["Prioritize team communication", "Check weather updates"]
+            quote = "Every step forward counts!"
+            response_json = {"checklist": checklist, "tips": tips, "quote": quote}
+            logger.warning("Model output is not JSON, using default response")
+        else:
+            try:
+                response_json = json.loads(response_text)
+            except json.JSONDecodeError:
+                response_json = {
+                    "checklist": ["Inspect safety equipment", "Review milestone progress"],
+                    "tips": ["Prioritize team communication", "Check weather updates"],
+                    "quote": "Every step forward counts!"
+                }
+                logger.warning("Failed to parse model output as JSON, using default response")
         
         return response_json