Update README.md

README.md

@@ -34,26 +34,94 @@ The model can predict the following emotions in text:
 Here is an example of how to run inference with the model:
 
 ```python
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
+from torch import nn
+from transformers import AutoTokenizer, AutoModel
+import safetensors.torch
+import pandas as pd
+
+# Custom model class for emotion classification using MiniLM
+class MiniLMEmotionClassifier(nn.Module):
+    def __init__(self, model_name):
+        super(MiniLMEmotionClassifier, self).__init__()
+        self.base_model = AutoModel.from_pretrained(model_name, ignore_mismatched_sizes=True)  # Load the MiniLM model
+        self.dropout = nn.Dropout(0.1)  # Dropout for regularization
+        self.fc = nn.Linear(384, 6)  # Output layer for 6 emotion categories
+
+    def forward(self, input_ids, attention_mask=None, labels=None):
+        outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.last_hidden_state[:, 0, :]  # Extract [CLS] token representation
+        pooled_output = self.dropout(pooled_output)
+        logits = self.fc(pooled_output)  # Compute predictions
+
+        loss = None
+        if labels is not None:
+            # Use MSE loss for regression-style emotion prediction
+            loss_fct = nn.MSELoss()
+            loss = loss_fct(logits, labels.view_as(logits))
+
+        return {"loss": loss, "logits": logits} if loss is not None else {"logits": logits}
+
+# Path to the safetensors model file
+model_path = 'MiniLM-L12-Affect/model.safetensors'
+
+# Load model weights from the safetensors file
+with open(model_path, 'rb') as f:
+    model_data = f.read()
+model_state_dict = safetensors.torch.load(model_data)
+
+# Initialize the MiniLM model
+model_name = "./MiniLM-L12-Affect"
+model = MiniLMEmotionClassifier(model_name)
+
+# Load pre-trained weights into the model
+model.load_state_dict(model_state_dict, strict = False)
+
+# Load the tokenizer
+tokenizer = AutoTokenizer.from_pretrained("./MiniLM-L12-Affect", ignore_mismatched_sizes=True)
+
+# Move model to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+model.eval()
 
-# Load the model and tokenizer
-model = AutoModelForSequenceClassification.from_pretrained("pharci/MiniLM-L12-Affect")
-tokenizer = AutoTokenizer.from_pretrained("pharci/MiniLM-L12-Affect")
-
-# Emotion prediction function
 def predict_emotions(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    """Tokenizes input text and predicts emotion scores."""
+    inputs = tokenizer(
+        text,
+        padding="max_length",
+        truncation=True,
+        max_length=128,
+        return_tensors="pt"
+    )
+    # Remove 'token_type_ids' if present
+    inputs.pop('token_type_ids', None)
+    inputs = {key: value.to(device) for key, value in inputs.items()}
+
     with torch.no_grad():
         outputs = model(**inputs)
-    predictions = outputs.logits.squeeze().cpu().numpy()
+
+    predictions = outputs["logits"].cpu().numpy()[0]
     return predictions
 
-# Example prediction
-test_text = "I am really happy today!"
-predictions = predict_emotions(test_text)
-print(predictions)
+# Example inference
+test_text = "This is horribly amazing ! you're a genius"
+emotions = predict_emotions(test_text)
+
+# Emotion categories
+categories = ["Joy", "Anger", "Fear", "Sadness", "Surprise", "Disgust"]
+
+# Display the results
+print(f"Text: {test_text}")
+emotion_df = pd.DataFrame(emotions.reshape(1, -1), columns=categories)
+print(emotion_df)
 ```
+**Result**
+|     | Joy     | Anger   | Fear    | Sadness  | Surprise | Disgust  |
+|-----|---------|--------|---------|---------|----------|---------|
+|  0  | 0.844805 | 0.02971 | 0.008245 | -0.007872 | 0.668609  | 0.001267 |
+
+
 
 ## Deployment