Update app.py

app.py

@@ -1,9 +1,11 @@
-from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from contextlib import asynccontextmanager
 import json
 import logging
 import os
+import asyncio
 
 # Set up logging to stdout only
 logging.basicConfig(
@@ -15,16 +17,6 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Initialize FastAPI app
-app = FastAPI()
-
-# Define input model for validation
-class CoachingInput(BaseModel):
-    role: str
-    project_id: str
-    milestones: str
-    reflection_log: str
-
 # Global variables for model and tokenizer
 model = None
 tokenizer = None
@@ -53,11 +45,24 @@ async def load_model_background():
         logger.error(f"Failed to load model or tokenizer: {str(e)}")
         model_load_status = f"failed: {str(e)}"
 
-# Startup event to initiate model loading in the background
-@app.on_event("startup")
-async def startup_event(background_tasks: BackgroundTasks):
-    logger.debug("FastAPI application started")
-    background_tasks.add_task(load_model_background)
+# Lifespan event handler to manage startup and shutdown
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    logger.debug("FastAPI application starting")
+    # Start the background task for model loading
+    asyncio.create_task(load_model_background())
+    yield
+    logger.debug("FastAPI application shutting down")
+
+# Initialize FastAPI app with lifespan handler
+app = FastAPI(lifespan=lifespan)
+
+# Define input model for validation
+class CoachingInput(BaseModel):
+    role: str
+    project_id: str
+    milestones: str
+    reflection_log: str
 
 @app.get("/")
 async def root():