Update app.py

app.py

@@ -1,12 +1,20 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, BackgroundTasks
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import json
 import logging
 import os
+import asyncio
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
+# Set up logging to both stdout and a file
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    handlers=[
+        logging.StreamHandler(),  # Log to stdout
+        logging.FileHandler("/app/app.log")  # Log to a file
+    ]
+)
 logger = logging.getLogger(__name__)
 
 # Initialize FastAPI app
@@ -26,10 +34,10 @@ model_load_status = "not_loaded"
 
 # Define model path and fallback
 model_path = "/app/fine-tuned-construction-llm"
-fallback_model = "distilgpt2"  # Smaller model for faster loading
+fallback_model = "distilgpt2"
 
-# Load model and tokenizer at startup
-def load_model():
+# Asynchronous function to load model in the background
+async def load_model_background():
     global model, tokenizer, model_load_status
     try:
         if os.path.isdir(model_path):
@@ -38,7 +46,7 @@ def load_model():
             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}")
+            logger.info(f"Model directory not found: {model_path}. Using pre-trained model: {fallback_model}")
             model = AutoModelForCausalLM.from_pretrained(fallback_model)
             tokenizer = AutoTokenizer.from_pretrained(fallback_model)
             model_load_status = "fallback_model_loaded"
@@ -46,24 +54,25 @@ def load_model():
     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()
 
+# Startup event to initiate model loading in the background
 @app.on_event("startup")
-async def startup_event():
+async def startup_event(background_tasks: BackgroundTasks):
     logger.info("FastAPI application started")
+    background_tasks.add_task(load_model_background)
 
 @app.get("/health")
 async def health_check():
-    return {"status": "healthy", "model_load_status": model_load_status}
+    return {
+        "status": "healthy" if model_load_status in ["local_model_loaded", "fallback_model_loaded"] else "starting",
+        "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.")
+        raise HTTPException(status_code=503, detail="Model not loaded yet. Please try again later.")
 
     try:
         # Prepare input text