BERT Model to classify gym sentiment

app.py

@@ -63,15 +63,23 @@ class MultiHeadBERT(nn.Module):
         
         return workout_logits, feeling_logits, soreness_logits
         
-class ClassificationRequest(BaseModel):
-    message: str
+class PredictRequest(BaseModel):
+    user_input: str
+
+class PredictResponse(BaseModel):
+    workout:        str
+    workout_conf:   float
+    feeling:        str
+    feeling_conf:   float
+    soreness:       str
+    soreness_conf:  float
     
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}
 
-@app.post("/classify")
-def sentiment_analysis(payload: ClassificationRequest):
+@app.post("/predict",response_model=PredictResponse)
+def predict(request: PredictRequest):
     
     model = MultiHeadBERT(
         num_workout_types=4,
@@ -86,14 +94,42 @@ def sentiment_analysis(payload: ClassificationRequest):
     model.to(device)
     model.eval()
 
-    result = predict(
-        text=payload.message,
-        model=model,
-        tokenizer=tokenizer,
-        device=device
+    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, feeling_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)
+    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)
+    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_soreness = soreness_label_map[soreness_logits.argmax().item()]
+
     
-    return {
-        result
-    }
+    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),
+    )