Update app.py

app.py

@@ -18,7 +18,7 @@ workout_label_map = {
     3: "HIIT"
 }
 
-feeling_label_map = {
+mood_label_map = {
     0: "Energized",
     1: "Tired",
     2: "Stressed",
@@ -32,7 +32,7 @@ soreness_label_map = {
 }
 
 class MultiHeadBERT(nn.Module):
-    def __init__(self, num_workout_types, num_feelings, num_soreness_levels):
+    def __init__(self, num_workout_types, num_moods, num_soreness_levels):
         super(MultiHeadBERT, self).__init__()
         
         # Shared BERT backbone
@@ -41,7 +41,7 @@ class MultiHeadBERT(nn.Module):
         
         # Task-specific classification heads
         self.workout_head = nn.Linear(hidden_size, num_workout_types)
-        self.feeling_head = nn.Linear(hidden_size, num_feelings)
+        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)
@@ -58,21 +58,21 @@ class MultiHeadBERT(nn.Module):
         
         # Each head produces its own logits
         workout_logits = self.workout_head(cls_output)
-        feeling_logits = self.feeling_head(cls_output)
+        mood_logits = self.mood_head(cls_output)
         soreness_logits = self.soreness_head(cls_output)
         
-        return workout_logits, feeling_logits, soreness_logits
+        return workout_logits, mood_logits, soreness_logits
         
 class PredictRequest(BaseModel):
     user_input: str
 
 class PredictResponse(BaseModel):
-    workout:        str
-    workout_conf:   float
-    feeling:        str
-    feeling_conf:   float
-    soreness:       str
-    soreness_conf:  float
+    workout: str
+    workout_conf: float
+    mood: str
+    mood_conf: float
+    soreness: str
+    soreness_conf: float
     
 @app.get("/")
 def greet_json():
@@ -82,8 +82,8 @@ def greet_json():
 def predict(request: PredictRequest):
     
     model = MultiHeadBERT(
-        num_workout_types=4,
-        num_feelings=4,
+        num_workout_types=8,
+        num_moods=5,
         num_soreness_levels=3
     )
     
@@ -104,34 +104,34 @@ def predict(request: PredictRequest):
         return_tensors='pt'
     )
 
-    input_ids      = encoding['input_ids'].to(device)
+    input_ids = encoding['input_ids'].to(device)
     attention_mask = encoding['attention_mask'].to(device)
 
     with torch.no_grad():
-        workout_logits, feeling_logits, soreness_logits = model(input_ids, attention_mask)
+        workout_logits, mood_logits, soreness_logits = model(input_ids, attention_mask)
         
     # Convert logits to probabilities
     workout_probs  = torch.softmax(workout_logits,  dim=1)
-    feeling_probs  = torch.softmax(feeling_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)
-    feeling_conf,  feeling_pred  = feeling_probs.max(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 predictions to labels
-    predicted_workout  = workout_label_map[workout_logits.argmax().item()]
-    predicted_feeling  = feeling_label_map[feeling_logits.argmax().item()]
+    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()]
 
     
     return PredictResponse(
-        workout        = predicted_workout,
-        workout_conf   = round(workout_conf.item()  * 100, 1),
-        feeling        = predicted_feeling,
-        feeling_conf   = round(feeling_conf.item()  * 100, 1),
-        soreness       = predicted_soreness,
-        soreness_conf  = round(soreness_conf.item() * 100, 1),
+        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),
     )