Update app.py

app.py

@@ -4,12 +4,12 @@ from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from typing import List 
+from typing import List
 
-from supabase import create_client, Client
 import torch
 import torch.nn as nn
 from transformers import DistilBertModel, DistilBertTokenizer
+from supabase import create_client, Client
 
 # ── Logging setup ─────────────────────────────────────────────────────────────
 logging.basicConfig(level=logging.INFO)
@@ -30,29 +30,24 @@ soreness_label_map = {
     0: "None",      1: "Mild",      2: "Severe"
 }
 
+# ── Model Definition ──────────────────────────────────────────────────────────
 class MultiHeadDistilBERT(nn.Module):
     def __init__(self, num_workout_types, num_moods, num_soreness_levels):
         super(MultiHeadDistilBERT, self).__init__()
-        
-        # Shared BERT backbone
-        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.bert        = DistilBertModel.from_pretrained(
+                               'distilbert-base-uncased',
+                               token=os.getenv('HF_TOKEN')
+                           )
+        hidden_size      = self.bert.config.hidden_size
+        self.dropout     = nn.Dropout(0.3)
+        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)
-        
 
     def forward(self, input_ids, attention_mask):
-        outputs = self.bert(input_ids=input_ids,attention_mask=attention_mask)
-        
-        # Use [CLS] token representation. DistilBERT uses last_hidden_state instead of pooler_output like BERT
-        cls_output = self.dropout(outputs.last_hidden_state[:, 0, :])  # [CLS] token is first token in sequence
-        
-        # Each head produces its own logits        
+        outputs    = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        cls_output = self.dropout(outputs.last_hidden_state[:, 0, :])
         return (
             self.workout_head(cls_output),
             self.mood_head(cls_output),
@@ -72,7 +67,6 @@ state = AppState()
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # ── Startup ───────────────────────────────────────────────────────────────
-
     logger.info("Loading model, tokenizer and Supabase client...")
 
     state.device = torch.device('cpu')
@@ -82,7 +76,6 @@ async def lifespan(app: FastAPI):
         'distilbert-base-uncased',
         token=os.getenv('HF_TOKEN')
     )
-
     logger.info("Tokenizer loaded")
 
     # Load model once
@@ -111,9 +104,9 @@ async def lifespan(app: FastAPI):
     # ── Shutdown ──────────────────────────────────────────────────────────────
     logger.info("Shutting down API")
 
-
 app = FastAPI(lifespan=lifespan)
 
+# ── Schemas ───────────────────────────────────────────────────────────────────
 class PredictRequest(BaseModel):
     user_input: str
 
@@ -125,16 +118,15 @@ class ExerciseResponse(BaseModel):
     notes:             str
     suitable_moods:    List[int]
     suitable_soreness: List[int]
-    
+
 class PredictResponse(BaseModel):
-    workout: str
-    workout_conf: float
-    mood: str
-    mood_conf: float
-    soreness: str
+    workout:       str
+    workout_conf:  float
+    mood:          str
+    mood_conf:     float
+    soreness:      str
     soreness_conf: float
-    exercises: List[ExerciseResponse]
-
+    exercises:     List[ExerciseResponse]
 
 # ── Supabase Helper ───────────────────────────────────────────────────────────
 def get_suitable_exercises(workout_type: int, mood: int, soreness: int) -> List[ExerciseResponse]:
@@ -152,7 +144,7 @@ def get_suitable_exercises(workout_type: int, mood: int, soreness: int) -> List[
     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 health_check():
@@ -160,7 +152,7 @@ def health_check():
         "status":  "ok",
         "model":   "MultiHeadDistilBERT",
         "device":  str(state.device)
-}
+    }
 
 # ── Predict Endpoint ──────────────────────────────────────────────────────────
 @app.post("/predict", response_model=PredictResponse)
@@ -232,4 +224,4 @@ def predict(request: PredictRequest):
     except Exception as e:
         logger.error(f"Prediction failed: {e}")
         raise HTTPException(status_code=500, detail="Prediction failed. Please try again.")
-    
+