Update README.md

README.md

@@ -67,14 +67,58 @@ For simplicity, you can run this here:
 Alternatively, embed the following code in your application:
 
 ```python
-from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer
+import torch
+
+from transformers import AutoTokenizer, AutoConfig,XLMRobertaForSequenceClassification, PreTrainedModel
+from torch import nn
+from torch.nn import Dropout
+
+
+# Define the CustomModel class which is predicting Both SENTIMENT POLARITY &  EMOTIONS
+class CustomModel(XLMRobertaForSequenceClassification):
+    def __init__(self, config, num_emotion_labels):
+        super(CustomModel, self).__init__(config)
+        self.num_emotion_labels = num_emotion_labels
+        self.dropout_emotion = nn.Dropout(config.hidden_dropout_prob)
+        self.emotion_classifier = nn.Sequential(
+            nn.Linear(config.hidden_size, 512),
+            nn.Mish(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_emotion_labels)
+        )
+        self._init_weights(self.emotion_classifier[0])
+        self._init_weights(self.emotion_classifier[3])
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+    def forward(self, input_ids=None, attention_mask=None, sentiment=None, labels=None):
+        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
+        sequence_output = outputs[0]
+        if len(sequence_output.shape) != 3:
+            raise ValueError(f"Expected sequence_output to have 3 dimensions, got {sequence_output.shape}")
+        cls_hidden_states = sequence_output[:, 0, :]
+        cls_hidden_states = self.dropout_emotion(cls_hidden_states)
+        emotion_logits = self.emotion_classifier(cls_hidden_states)
+        with torch.no_grad():
+            cls_token_state = sequence_output[:, 0, :].unsqueeze(1)
+            sentiment_logits = self.classifier(cls_token_state).squeeze(1)
+        if labels is not None:
+            class_weights = torch.tensor([1.0] * self.num_emotion_labels).to(labels.device)
+            loss_fct = nn.BCEWithLogitsLoss(pos_weight=class_weights)
+            loss = loss_fct(emotion_logits, labels)
+            return {"loss": loss, "emotion_logits": emotion_logits, "sentiment_logits": sentiment_logits}
+        return {"emotion_logits": emotion_logits, "sentiment_logits": sentiment_logits}
+
+
+# Load the tokenizer and model from the local directory
+model_dir = "gsar78/HellenicSentimentAI_v2"
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+config = AutoConfig.from_pretrained(model_dir)
+model = CustomModel.from_pretrained(model_dir, config=config, num_emotion_labels=18)
 
-model_name = "gsar78/HellenicSentimentAI_v2"
 
-# Load the configuration, tokenizer, and model
-config = AutoConfig.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
 # Function to predict sentiment and emotion
 def predict(texts):
@@ -131,6 +175,8 @@ def predict(texts):
 sample_texts = ["Απολαύσαμε μια υπέροχη βραδιά σε αυτό το εστιατόριο. "
 "Το μενού ήταν πολύ καλά σχεδιασμένο και κάθε πιάτο ήταν μια γευστική έκπληξη. "
 "Η εξυπηρέτηση ήταν άψογη και η ατμόσφαιρα ευχάριστη. Σίγουρα θα επιστρέψουμε για άλλη μια φορά."]
+
+
 print("Text: ", sample_texts[0])
 emotion_results, sentiment_results = predict(sample_texts)
 
@@ -147,10 +193,16 @@ for label, prob in emotion_results.items():
 # Change the text and predict again
 # Print the results
 print("\n======")
+
+
 print("\nNew prediction:")
 sample_texts = ["Η τελευταία μας εμπειρία στο εστιατόριο αυτό δεν ήταν ιδιαίτερα θετική. "
 "Αν και ο χώρος είχε μια ενδιαφέρουσα ατμόσφαιρα, το φαγητό ήταν μέτριο και η εξυπηρέτηση ήταν αργή. "
 "Οι τιμές ήταν επίσης απογοητευτικές για την ποιότητα που προσφέρθηκε."]
+
+
+
+
 print("Text: ", sample_texts[0])
 emotion_results, sentiment_results = predict(sample_texts)
 
@@ -160,6 +212,7 @@ for label, prob in sentiment_results.items():
 print("\nEmotion probabilities (%):")
 for label, prob in emotion_results.items():
     print(f"    {label}: {prob:.2f}%")
+
 ```
 
 Expected output: