Changed to DistilBERT model

app.py

@@ -5,7 +5,7 @@ from pydantic import BaseModel
 
 import torch
 import torch.nn as nn
-from transformers import BertModel, BertTokenizer
+from transformers import DistilBertModel, DistilBertTokenizer
 
 app = FastAPI()
 
@@ -36,12 +36,12 @@ soreness_label_map = {
     2: "Severe"
 }
 
-class MultiHeadBERT(nn.Module):
+class MultiHeadDistilBERT(nn.Module):
     def __init__(self, num_workout_types, num_moods, num_soreness_levels):
-        super(MultiHeadBERT, self).__init__()
+        super(MultiHeadDistilBERT, self).__init__()
         
         # Shared BERT backbone
-        self.bert = BertModel.from_pretrained('bert-base-uncased',token=os.getenv('HF_TOKEN'))
+        self.bert = DistilBertModel.from_pretrained('distilbert-base-uncased',token=os.getenv('HF_TOKEN'))
         hidden_size = self.bert.config.hidden_size  # 768
         
         # Task-specific classification heads
@@ -51,22 +51,14 @@ class MultiHeadBERT(nn.Module):
         
         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
-        )
+    def forward(self, input_ids, attention_mask):
+        outputs = self.bert(input_ids=input_ids,attention_mask=attention_mask)
         
-        # Use [CLS] token representation
-        cls_output = self.dropout(outputs.pooler_output)
+        # 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
-        workout_logits = self.workout_head(cls_output)
-        mood_logits = self.mood_head(cls_output)
-        soreness_logits = self.soreness_head(cls_output)
-        
-        return workout_logits, mood_logits, soreness_logits
+        # Each head produces its own logits        
+        return  (self.workout_head(cls_output), self.mood_head(cls_output), self.soreness_head(cls_output))     
         
 class PredictRequest(BaseModel):
     user_input: str
@@ -86,20 +78,17 @@ def greet_json():
 @app.post("/predict",response_model=PredictResponse)
 def predict(request: PredictRequest):
     
-    model = MultiHeadBERT(
+    model = MultiHeadDistilBERT(
         num_workout_types=8,
         num_moods=5,
         num_soreness_levels=3
     )
     
-    model.load_state_dict(
-        torch.load('best_model.pt', map_location=torch.device('cpu'))
-    )
-    
+    model.load_state_dict(torch.load('best_model.pt', map_location=torch.device('cpu')))
     model.to(device)
     model.eval()
 
-    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased',token=os.getenv('HF_TOKEN'))
+    tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased',token=os.getenv('HF_TOKEN'))
 
     encoding = tokenizer(
         request.user_input,  # The single string the user types