Enhance subjectivity prediction with detailed output and update Gradio interface

app.py

@@ -7,7 +7,8 @@ import torch.nn as nn
 
 # Define the model and tokenizer
 model_card = "microsoft/mdeberta-v3-base"
-finetuned_model = "MatteoFasulo/mdeberta-v3-base-subjectivity-sentiment-multilingual"
+finetuned_model = "MatteoFasulo/mdeberta-v3-base-subjectivity-sentiment-multilingual-no-arabic"
+THRESHOLD = 0.65
 
 # Custom model class for combining sentiment analysis with subjectivity detection
 class CustomModel(PreTrainedModel):
@@ -22,7 +23,7 @@ class CustomModel(PreTrainedModel):
 
         self.classifier = nn.Linear(output_dim + sentiment_dim, num_labels)
 
-    def forward(self, input_ids, positive, neutral, negative, attention_mask=None, labels=None):
+    def forward(self, input_ids, positive, neutral, negative, token_type_ids=None, attention_mask=None, labels=None):
         outputs = self.deberta(input_ids=input_ids, attention_mask=attention_mask)
 
         encoder_layer = outputs[0]
@@ -66,24 +67,48 @@ def get_sentiment_values(text: str):
     sentiments = pipe(text)[0]
     return {k:v for k,v in [(list(sentiment.values())[0], list(sentiment.values())[1]) for sentiment in sentiments]}
 
-# Predict the subjectivity of a sentence
+# Modify the predict_subjectivity function to return additional information
 def predict_subjectivity(text):
     sentiment_values = get_sentiment_values(text)
 
     model = load_model(model_card, finetuned_model)
     tokenizer = load_tokenizer(model_card)
 
+    positive = sentiment_values['positive']
+    neutral = sentiment_values['neutral']
+    negative = sentiment_values['negative']
+
     inputs = tokenizer(text, padding=True, truncation=True, max_length=256, return_tensors='pt')
+    inputs['positive'] = torch.tensor(positive).unsqueeze(0)
+    inputs['neutral'] = torch.tensor(neutral).unsqueeze(0)
+    inputs['negative'] = torch.tensor(negative).unsqueeze(0)
 
     outputs = model(**inputs)
     logits = outputs.get('logits')
 
-    predicted_class_idx = logits.argmax().item()
+    # Calculate probabilities using softmax
+    probabilities = torch.nn.functional.softmax(logits, dim=1)
+    obj_prob, subj_prob = probabilities[0].tolist()
+
+    # Predict the class given the decision threshold
+    predicted_class_idx = 1 if subj_prob >= THRESHOLD else 0
     predicted_class = model.config.id2label[predicted_class_idx]
 
-    return predicted_class
+    # Format the output
+    result = f"""Prediction: {predicted_class}
+
+Class Probabilities:
+- Objective: {obj_prob:.2%}
+- Subjective: {subj_prob:.2%}
+
+Sentiment Scores:
+- Positive: {positive:.2%}
+- Neutral: {neutral:.2%}
+- Negative: {negative:.2%}"""
+
+    return result
 
-# Create a Gradio interface
+# Update the Gradio interface
 demo = gr.Interface(
     fn=predict_subjectivity, 
     inputs=gr.Textbox(
@@ -91,14 +116,12 @@ demo = gr.Interface(
         placeholder='Enter a sentence from a news article',
         info='Paste a sentence from a news article to determine if it is subjective or objective.'
     ),
-    outputs=gr.Text(
-        label="Prediction",
-        info="Whether the sentence is subjective or objective."
+    outputs=gr.Textbox(
+        label="Results",
+        info="Detailed analysis including subjectivity prediction, class probabilities, and sentiment scores."
     ),
     title='Subjectivity Detection',
-    description='Detect if a sentence is subjective or objective using a pre-trained model.',
-    theme='huggingface',
+    description='Detect if a sentence is subjective or objective using a pre-trained model.'
 )
 
-# Launch the interface
-demo.launch(share=True)
+demo.launch()