Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py
 import gradio as gr
 import torch
 from transformers import (
@@ -15,7 +14,7 @@ DEVICE = "cpu"  # HF Spaces обычно CPU
 # =====================
 # Agreement (MNLI)
 # =====================
-MNLI_MODEL = "facebook/bart-base-mnli"
+MNLI_MODEL = "facebook/bart-large-mnli"  # Fixed: valid public model
 mnli_tokenizer = None
 mnli_model = None
 
@@ -33,6 +32,7 @@ def check_agreement(msg1: str, msg2: str) -> float:
     with torch.no_grad():
         logits = mnli_model(**inputs).logits
     probs = torch.softmax(logits, dim=-1)[0]
+    # Agreement score: entailment - contradiction
     return round((probs[2] - probs[0]).item(), 2)
 
 # =====================
@@ -57,6 +57,7 @@ def analyze_sentiment(text: str) -> float:
         logits = sent_model(**inputs).logits
     probs = torch.softmax(logits, dim=-1)
     stars = torch.argmax(probs, dim=-1).item() + 1
+    # Convert 1-5 stars into -5 to +5 scale
     return round((stars - 3) * 2.5, 2)
 
 # =====================
@@ -82,14 +83,14 @@ def load_classifier():
         clf_model.to(DEVICE)
         clf_model.eval()
 
-def classify_message(text: str) -> list:
+def classify_message(text: str) -> dict:
     load_classifier()
     inputs = clf_tokenizer(text, return_tensors="pt", truncation=True).to(DEVICE)
     with torch.no_grad():
         logits = clf_model(**inputs).logits
     probs = torch.sigmoid(logits)[0]
-    labels = [CATEGORIES[i] for i, p in enumerate(probs) if p > 0.5]
-    return labels or ["neutral"]
+    # Return a dict for Gradio Label output
+    return {CATEGORIES[i]: float(probs[i]) for i in range(len(CATEGORIES))}
 
 # =====================
 # Gradio UI
@@ -97,6 +98,7 @@ def classify_message(text: str) -> list:
 with gr.Blocks(title="Unified NLP API") as demo:
     gr.Markdown("## 📈 Unified NLP API")
     
+    # ----- Agreement Tab -----
     with gr.Tab("Agreement"):
         msg1 = gr.Textbox(label="Message 1")
         msg2 = gr.Textbox(label="Message 2")
@@ -104,12 +106,14 @@ with gr.Blocks(title="Unified NLP API") as demo:
         out_agree = gr.Number(label="Agreement Score")
         btn_agree.click(fn=check_agreement, inputs=[msg1, msg2], outputs=out_agree)
     
+    # ----- Sentiment Tab -----
     with gr.Tab("Sentiment"):
         text_sent = gr.Textbox(label="Text")
         btn_sent = gr.Button("Analyze Sentiment")
         out_sent = gr.Number(label="Sentiment Score (-5 to +5)")
         btn_sent.click(fn=analyze_sentiment, inputs=text_sent, outputs=out_sent)
     
+    # ----- Multilabel Classification Tab -----
     with gr.Tab("Multilabel Classification"):
         text_clf = gr.Textbox(label="Text")
         btn_clf = gr.Button("Classify")