Update app.py

app.py

@@ -9,6 +9,7 @@ import os
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Initialize FastAPI app
 app = FastAPI()
 
 # Define input model for validation
@@ -18,27 +19,52 @@ class CoachingInput(BaseModel):
     milestones: str
     reflection_log: str
 
-# Define model path (absolute path in the container)
+# Global variables for model and tokenizer
+model = None
+tokenizer = None
+model_load_status = "not_loaded"
+
+# Define model path and fallback
 model_path = "/app/fine-tuned-construction-llm"
-fallback_model = "gpt2"  # Fallback to a pre-trained model if local model is unavailable
+fallback_model = "distilgpt2"  # Smaller model for faster loading
 
-# Load model and tokenizer
-try:
-    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)}")
-    raise Exception(f"Model loading failed: {str(e)}")
+# Load model and tokenizer at startup
+def load_model():
+    global model, tokenizer, model_load_status
+    try:
+        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)
+            model_load_status = "local_model_loaded"
+        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)
+            model_load_status = "fallback_model_loaded"
+        logger.info("Model and tokenizer loaded successfully")
+    except Exception as e:
+        logger.error(f"Failed to load model or tokenizer: {str(e)}")
+        model_load_status = f"failed: {str(e)}"
+        # Do not raise an exception; allow the app to start
+
+# Load model on startup
+load_model()
+
+@app.on_event("startup")
+async def startup_event():
+    logger.info("FastAPI application started")
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy", "model_load_status": model_load_status}
 
 @app.post("/generate_coaching")
 async def generate_coaching(data: CoachingInput):
+    if model is None or tokenizer is None:
+        logger.error("Model or tokenizer not loaded")
+        raise HTTPException(status_code=503, detail="Model not loaded. Please check server logs.")
+
     try:
         # Prepare input text
         input_text = (
@@ -62,8 +88,7 @@ async def generate_coaching(data: CoachingInput):
         # Decode and parse response
         response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # 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
+        # Since distilgpt2 may not output JSON, parse the response manually or use fallback
         if not response_text.startswith("{"):
             checklist = ["Inspect safety equipment", "Review milestone progress"]
             tips = ["Prioritize team communication", "Check weather updates"]
@@ -85,8 +110,4 @@ async def generate_coaching(data: CoachingInput):
     
     except Exception as e:
         logger.error(f"Error generating coaching response: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
-
-@app.get("/health")
-async def health_check():
-    return {"status": "healthy"}
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")