Update app.py

app.py

@@ -9,33 +9,23 @@ import torch
 import torch.nn as nn
 from transformers import DistilBertModel, DistilBertTokenizer
 
-app = FastAPI()
-
-device = torch.device('cpu')  # Hugging Face Space with no GPU
+# ── Logging setup ─────────────────────────────────────────────────────────────
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
+# ── Label Maps ────────────────────────────────────────────────────────────────
 workout_label_map = {
-    0: "Chest",
-    1: "Back",
-    2: "Legs",
-    3: "Shoulders",
-    4: "Arms",
-    5: "Core",
-    6: "Full Body",
-    7: "Cardio"
+    0: "Chest",     1: "Back",      2: "Legs",      3: "Shoulders",
+    4: "Arms",      5: "Core",      6: "Full Body",  7: "Cardio"
 }
 
 mood_label_map = {
-    0: "Energized",
-    1: "Tired",
-    2: "Stressed",
-    3: "Motivated",
-    4: "Neutral"
+    0: "Energized", 1: "Tired",     2: "Stressed",
+    3: "Motivated", 4: "Neutral"
 }
 
 soreness_label_map = {
-    0: "None",
-    1: "Mild",
-    2: "Severe"
+    0: "None",      1: "Mild",      2: "Severe"
 }
 
 class MultiHeadDistilBERT(nn.Module):
@@ -46,12 +36,13 @@ class MultiHeadDistilBERT(nn.Module):
         self.bert = DistilBertModel.from_pretrained('distilbert-base-uncased',token=os.getenv('HF_TOKEN'))
         hidden_size = self.bert.config.hidden_size  # 768
         
+        self.dropout = nn.Dropout(0.3)
+
         # Task-specific classification heads
         self.workout_head = nn.Linear(hidden_size, num_workout_types)
         self.mood_head = nn.Linear(hidden_size, num_moods)
         self.soreness_head = nn.Linear(hidden_size, num_soreness_levels)
         
-        self.dropout = nn.Dropout(0.3)
 
     def forward(self, input_ids, attention_mask):
         outputs = self.bert(input_ids=input_ids,attention_mask=attention_mask)
@@ -60,8 +51,65 @@ class MultiHeadDistilBERT(nn.Module):
         cls_output = self.dropout(outputs.last_hidden_state[:, 0, :])  # [CLS] token is first token in sequence
         
         # Each head produces its own logits        
-        return  (self.workout_head(cls_output), self.mood_head(cls_output), self.soreness_head(cls_output))     
-        
+        return (
+            self.workout_head(cls_output),
+            self.mood_head(cls_output),
+            self.soreness_head(cls_output)
+        )
+
+# ── App State — loaded once at startup ───────────────────────────────────────
+class AppState:
+    model:     MultiHeadDistilBERT = None
+    tokenizer: DistilBertTokenizer = None
+    supabase:  Client              = None
+    device:    torch.device        = None
+
+state = AppState()
+
+# ── Lifespan — runs once on startup and shutdown ──────────────────────────────
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # ── Startup ───────────────────────────────────────────────────────────────
+    logger.info("Loading model, tokenizer and Supabase client...")
+
+    state.device = torch.device('cpu')
+
+    # Load tokenizer once
+    state.tokenizer = DistilBertTokenizer.from_pretrained(
+        'distilbert-base-uncased',
+        token=os.getenv('HF_TOKEN')
+    )
+    logger.info("Tokenizer loaded")
+
+    # Load model once
+    state.model = MultiHeadDistilBERT(
+        num_workout_types=8,
+        num_moods=5,
+        num_soreness_levels=3
+    )
+    state.model.load_state_dict(
+        torch.load('best_DistilBERT_model.pt', map_location=state.device)
+    )
+    state.model.to(state.device)
+    state.model.eval()
+    logger.info("Model loaded")
+
+    # Create Supabase client once
+    state.supabase = create_client(
+        os.getenv('SUPA_URL'),
+        os.getenv('SUPA_KEY')
+    )
+    logger.info("Supabase client created")
+    logger.info("Startup complete — API is ready")
+
+    yield  # ← API runs here
+
+    # ── Shutdown ──────────────────────────────────────────────────────────────
+    logger.info("Shutting down API")
+
+
+app = FastAPI(lifespan=lifespan)
+
 class PredictRequest(BaseModel):
     user_input: str
 
@@ -84,83 +132,100 @@ class PredictResponse(BaseModel):
     exercises: List[ExerciseResponse]
 
 
-def get_suitable_exercises(workout_type: int, mood: int, soreness: int):
-    supabase = create_client(os.getenv('SUPA_URL'), os.getenv('SUPA_KEY'))
-
-    supabase_response = (
-        supabase.table('exerciseai')
-        .select('*')
-        .eq('workout_type', workout_type)
-        .contains('suitable_moods',    [str(mood)])
-        .contains('suitable_soreness', [str(soreness)])
-        .execute()
-    )
-
-     # Parse Supabase response into ExerciseResponse objects
-    exercises = [ExerciseResponse(**exercise) for exercise in supabase_response.data]
-    return exercises
+# ── Supabase Helper ───────────────────────────────────────────────────────────
+def get_suitable_exercises(workout_type: int, mood: int, soreness: int) -> List[ExerciseResponse]:
+    try:
+        response = (
+            state.supabase.table('exerciseai')
+            .select('*')
+            .eq('workout_type', workout_type)
+            .contains('suitable_moods',    [mood])
+            .contains('suitable_soreness', [soreness])
+            .execute()
+        )
+        return [ExerciseResponse(**exercise) for exercise in response.data]
+
+    except Exception as e:
+        logger.error(f"Supabase query failed: {e}")
+        raise HTTPException(status_code=503, detail="Failed to fetch exercises from database")
     
+# ── Health Check ──────────────────────────────────────────────────────────────
 @app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+def health_check():
+    return {
+        "status":  "ok",
+        "model":   "MultiHeadDistilBERT",
+        "device":  str(state.device)
+}
 
-@app.post("/predict",response_model=PredictResponse)
+# ── Predict Endpoint ──────────────────────────────────────────────────────────
+@app.post("/predict", response_model=PredictResponse)
 def predict(request: PredictRequest):
-    
-    model = MultiHeadDistilBERT(
-        num_workout_types=8,
-        num_moods=5,
-        num_soreness_levels=3
-    )
-    
-    model.load_state_dict(torch.load('best_DistilBERT_model.pt', map_location=torch.device('cpu')))
-    model.to(device)
-    model.eval()
-
-    tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased',token=os.getenv('HF_TOKEN'))
-
-    encoding = tokenizer(
-        request.user_input,  # The single string the user types
-        max_length=128,
-        padding='max_length',
-        truncation=True,
-        return_tensors='pt'
-    )
-
-    input_ids = encoding['input_ids'].to(device)
-    attention_mask = encoding['attention_mask'].to(device)
 
-    with torch.no_grad():
-        workout_logits, mood_logits, soreness_logits = model(input_ids, attention_mask)
-        
-    # Convert logits to probabilities
-    workout_probs  = torch.softmax(workout_logits,  dim=1)
-    mood_probs  = torch.softmax(mood_logits,  dim=1)
-    soreness_probs = torch.softmax(soreness_logits, dim=1)
-
-    # Get predicted class and confidence percentage for each head
-    workout_conf, workout_pred  = workout_probs.max(dim=1)
-    mood_conf, mood_pred  = mood_probs.max(dim=1)
-    soreness_conf, soreness_pred = soreness_probs.max(dim=1)
-
-    # Map predictions to labels
-    predicted_workout = workout_label_map[workout_logits.argmax().item()]
-    predicted_mood = mood_label_map[mood_logits.argmax().item()]
-    predicted_soreness = soreness_label_map[soreness_logits.argmax().item()]
-
-    # Fetch suitable exercises from Supabase
-    suitable_exercises = get_suitable_exercises(
-        workout_type = workout_logits.argmax().item(),
-        mood         = mood_logits.argmax().item(),
-        soreness     = soreness_logits.argmax().item()
-    )    
-    return PredictResponse(
-        workout = predicted_workout,
-        workout_conf = round(workout_conf.item()  * 100, 1),
-        mood = predicted_mood,
-        mood_conf = round(mood_conf.item()  * 100, 1),
-        soreness = predicted_soreness,
-        soreness_conf = round(soreness_conf.item() * 100, 1),
-        exercises = suitable_exercises
-    )
+    # ── Input validation ──────────────────────────────────────────────────────
+    if not request.user_input.strip():
+        raise HTTPException(status_code=400, detail="user_input cannot be empty")
+
+    try:
+        # ── Tokenize ──────────────────────────────────────────────────────────
+        encoding = state.tokenizer(
+            request.user_input,
+            max_length=128,
+            padding='max_length',
+            truncation=True,
+            return_tensors='pt'
+        )
+
+        input_ids      = encoding['input_ids'].to(state.device)
+        attention_mask = encoding['attention_mask'].to(state.device)
+
+        # ── Inference ─────────────────────────────────────��───────────────────
+        with torch.no_grad():
+            workout_logits, mood_logits, soreness_logits = state.model(
+                input_ids, attention_mask
+            )
+
+        # ── Softmax + confidence ──────────────────────────────────────────────
+        workout_probs  = torch.softmax(workout_logits,  dim=1)
+        mood_probs     = torch.softmax(mood_logits,     dim=1)
+        soreness_probs = torch.softmax(soreness_logits, dim=1)
+
+        workout_conf,  workout_pred  = workout_probs.max(dim=1)
+        mood_conf,     mood_pred     = mood_probs.max(dim=1)
+        soreness_conf, soreness_pred = soreness_probs.max(dim=1)
+
+        # ── Map to labels — reuse pred variables, no redundant argmax ─────────
+        predicted_workout  = workout_label_map[workout_pred.item()]
+        predicted_mood     = mood_label_map[mood_pred.item()]
+        predicted_soreness = soreness_label_map[soreness_pred.item()]
+
+        logger.info(
+            f"Prediction — Workout: {predicted_workout} ({workout_conf.item()*100:.1f}%) | "
+            f"Mood: {predicted_mood} ({mood_conf.item()*100:.1f}%) | "
+            f"Soreness: {predicted_soreness} ({soreness_conf.item()*100:.1f}%)"
+        )
+
+        # ── Fetch exercises ───────────────────────────────────────────────────
+        suitable_exercises = get_suitable_exercises(
+            workout_type = workout_pred.item(),
+            mood         = mood_pred.item(),
+            soreness     = soreness_pred.item()
+        )
+
+        return PredictResponse(
+            workout       = predicted_workout,
+            workout_conf  = round(workout_conf.item()  * 100, 1),
+            mood          = predicted_mood,
+            mood_conf     = round(mood_conf.item()     * 100, 1),
+            soreness      = predicted_soreness,
+            soreness_conf = round(soreness_conf.item() * 100, 1),
+            exercises     = suitable_exercises
+        )
+
+    except HTTPException:
+        raise  # ← re-raise HTTP exceptions from get_suitable_exercises
+
+    except Exception as e:
+        logger.error(f"Prediction failed: {e}")
+        raise HTTPException(status_code=500, detail="Prediction failed. Please try again.")