Update app.py

app.py

@@ -1,12 +1,68 @@
+import os
+
 from fastapi import FastAPI
 from pydantic import BaseModel
 
-import joblib
-import pandas as pd
-import maven_text_preprocessing
+import torch
+import torch.nn as nn
+from transformers import BertModel, BertTokenizer
 
 app = FastAPI()
 
+device = torch.device('cpu')  # Hugging Face Space with no GPU
+
+workout_label_map = {
+    0: "Cardio",
+    1: "Strength",
+    2: "Yoga",
+    3: "HIIT"
+}
+
+feeling_label_map = {
+    0: "Energized",
+    1: "Tired",
+    2: "Stressed",
+    3: "Motivated"
+}
+
+soreness_label_map = {
+    0: "None",
+    1: "Mild",
+    2: "Severe"
+}
+
+class MultiHeadBERT(nn.Module):
+    def __init__(self, num_workout_types, num_feelings, num_soreness_levels):
+        super(MultiHeadBERT, self).__init__()
+        
+        # Shared BERT backbone
+        self.bert = BertModel.from_pretrained('bert-base-uncased',token=os.get_env('HF_TOKEN'))
+        hidden_size = self.bert.config.hidden_size  # 768
+        
+        # Task-specific classification heads
+        self.workout_head = nn.Linear(hidden_size, num_workout_types)
+        self.feeling_head = nn.Linear(hidden_size, num_feelings)
+        self.soreness_head = nn.Linear(hidden_size, num_soreness_levels)
+        
+        self.dropout = nn.Dropout(0.3)
+
+    def forward(self, input_ids, attention_mask, token_type_ids=None):
+        outputs = self.bert(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids
+        )
+        
+        # Use [CLS] token representation
+        cls_output = self.dropout(outputs.pooler_output)
+        
+        # Each head produces its own logits
+        workout_logits = self.workout_head(cls_output)
+        feeling_logits = self.feeling_head(cls_output)
+        soreness_logits = self.soreness_head(cls_output)
+        
+        return workout_logits, feeling_logits, soreness_logits
+        
 class ClassificationRequest(BaseModel):
     message: str
     
@@ -17,25 +73,27 @@ def greet_json():
 @app.post("/classify")
 def sentiment_analysis(payload: ClassificationRequest):
     
-    model = joblib.load("naive_bayes.joblib")
-    vectorizer = joblib.load("vectorizer.joblib")
-
-    clean_text = maven_text_preprocessing.clean_and_normalize(pd.Series([payload.message]))
-
-    X = vectorizer.transform(clean_text)  # ⚠️ transform, NOT fit_transform
-
-    category_list = ["Politics", "Sport", "Technology", "Entertainment", "Business"]
+    model = MultiHeadBERT(
+        num_workout_types=4,
+        num_feelings=4,
+        num_soreness_levels=3
+    )
     
-    predictions = model.predict(X)
+    model.load_state_dict(
+        torch.load('best_model.pt', map_location=torch.device('cpu'))
+    )
     
-    pred_prob = model.predict_proba(X)
-    pred_prob = pred_prob.tolist()[0]
+    model.to(device)
+    model.eval()
+
+    result = predict(
+        text=payload.message,
+        model=model,
+        tokenizer=tokenizer,
+        device=device
+    )
     
     return {
-        category_list[0]: pred_prob[0],
-        category_list[1]: pred_prob[1],
-        category_list[2]: pred_prob[2],
-        category_list[3]: pred_prob[3],
-        category_list[4]: pred_prob[4]
+        result
     }