Update app.py

app.py

@@ -1,12 +1,18 @@
 import torch
 import torch.nn as nn
-from transformers import BertTokenizer, BertModel
+from transformers import BertTokenizer, BertModel, BertForSequenceClassification
 import gradio as gr
 
-model_name = "keethu/bert-emotion-classifier"  
+# Your model repo
+model_name = "keethu/bert-emotion-classifier"
+
+# Load tokenizer
 tokenizer = BertTokenizer.from_pretrained(model_name)
-bert_model = BertModel.from_pretrained(model_name)
 
+# Load base BERT model
+base_bert = BertModel.from_pretrained(model_name)
+
+# Define your classifier architecture (same as training)
 class BERTClassifier(nn.Module):
     def __init__(self, bert_model, num_labels=5, dropout=0.3):
         super(BERTClassifier, self).__init__()
@@ -21,8 +27,19 @@ class BERTClassifier(nn.Module):
         logits = self.classifier(pooled_output)
         return logits
 
-model = BERTClassifier(bert_model, num_labels=5, dropout=0.3)
-model.load_state_dict(torch.load(f"{model_name}/pytorch_model.bin", map_location='cpu'))
+# Create model instance
+model = BERTClassifier(base_bert, num_labels=5, dropout=0.3)
+
+# Load the trained weights - USE from_pretrained properly
+from huggingface_hub import hf_hub_download
+import os
+
+# Download the model file
+model_path = hf_hub_download(repo_id=model_name, filename="pytorch_model.bin")
+
+# Load state dict
+state_dict = torch.load(model_path, map_location='cpu')
+model.load_state_dict(state_dict)
 model.eval()
 
 emotion_labels = ['anger', 'fear', 'joy', 'sadness', 'surprise']