Update app.py

app.py

@@ -12,7 +12,7 @@ from transformers import (
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # =====================
-# Agreement (MNLI)
+# 1) Agreement (MNLI)
 # =====================
 MNLI_MODEL = "facebook/bart-large-mnli"
 mnli_tokenizer = None
@@ -26,17 +26,28 @@ def load_mnli():
         mnli_model.to(DEVICE)
         mnli_model.eval()
 
-def check_agreement(msg1: str, msg2: str) -> float:
+def agreement_raw_score(msg1: str, msg2: str) -> float:
+    """
+    Возвращает "сырое" согласие в диапазоне [-1..+1]
+    по формуле entailment - contradiction.
+    """
     load_mnli()
     inputs = mnli_tokenizer(msg1, msg2, return_tensors="pt", truncation=True).to(DEVICE)
     with torch.no_grad():
         logits = mnli_model(**inputs).logits
     probs = torch.softmax(logits, dim=-1)[0]
-    # Считаем: entailment - contradiction
-    return round((probs[2] - probs[0]).item(), 2)
+    raw = (probs[2] - probs[0]).item()  # [-1..+1]
+    return raw
+
+def agreement_score_minus5_plus5(msg1: str, msg2: str) -> float:
+    """
+    Agreement в шкале [-5..+5]
+    """
+    raw = agreement_raw_score(msg1, msg2)
+    return round(raw * 5, 2)
 
 # =====================
-# Sentiment
+# 2) Sentiment (-5..+5)
 # =====================
 SENTIMENT_MODEL = "nlptown/bert-base-multilingual-uncased-sentiment"
 sent_tokenizer = None
@@ -51,17 +62,81 @@ def load_sentiment():
         sent_model.eval()
 
 def analyze_sentiment(text: str) -> float:
+    """
+    Модель даёт 1..5 звёзд -> переводим в [-5..+5]
+    """
     load_sentiment()
     inputs = sent_tokenizer(text, return_tensors="pt", truncation=True).to(DEVICE)
     with torch.no_grad():
         logits = sent_model(**inputs).logits
     probs = torch.softmax(logits, dim=-1)
-    stars = torch.argmax(probs, dim=-1).item() + 1
-    # Приводим шкалу 1–5 к -5..+5
-    return round((stars - 3) * 2.5, 2)
+    stars = torch.argmax(probs, dim=-1).item() + 1  # 1..5
+    score = (stars - 3) * 2.5  # -5..+5
+    return round(score, 2)
+
+# =====================
+# 3) Sarcasm / Irony (-5..+5)
+# =====================
+# Можно заменить модель на другую, если хочешь.
+# Эта модель популярна для сарказма.
+SARCASM_MODEL = "cardiffnlp/twitter-roberta-base-irony"
+sarcasm_pipe = None
+
+def load_sarcasm():
+    global sarcasm_pipe
+    if sarcasm_pipe is None:
+        sarcasm_pipe = pipeline(
+            "text-classification",
+            model=SARCASM_MODEL,
+            device=0 if torch.cuda.is_available() else -1,
+            truncation=True
+        )
+
+def sarcasm_score(text: str) -> float:
+    """
+    Возвращает рейтинг сарказма в [-5..+5]
+    (чем выше, тем больше сарказма/иронии)
+    """
+    load_sarcasm()
+    res = sarcasm_pipe(text)[0]
+    # Обычно метки: "irony" / "non_irony"
+    label = res["label"].lower()
+    conf = float(res["score"])  # 0..1
+
+    if "irony" in label:
+        # 0..1 -> 0..+5
+        return round(conf * 5, 2)
+    else:
+        # 0..1 -> 0..-5
+        return round(-conf * 5, 2)
+
+# =====================
+# 4) Agreement + Sarcasm
+# =====================
+def agreement_with_irony(msg1: str, msg2: str) -> float:
+    """
+    Идея:
+    - считаем agreement [-5..+5]
+    - считаем сарказм msg2 (обычно сарказм в ответе важнее)
+    - если сарказм высокий, уменьшаем "уверенность" agreement
+
+    Это НЕ идеальная логика, но работает лучше, чем игнорировать иронию.
+    """
+    base = agreement_score_minus5_plus5(msg1, msg2)
+
+    s2 = sarcasm_score(msg2)  # [-5..+5]
+    sarcasm_strength = abs(s2) / 5.0  # 0..1
+
+    # Чем больше сарказм, тем сильнее "сжимаем" agreement к нулю
+    # 0 сарказма -> множитель 1
+    # сильный сарказм -> множитель ~0.35
+    multiplier = 1.0 - 0.65 * sarcasm_strength
+
+    final_score = base * multiplier
+    return round(final_score, 2)
 
 # =====================
-# Zero‑Shot Classification
+# 5) Zero-Shot Multilabel -> [-5..+5]
 # =====================
 ZS_MODEL = "facebook/bart-large-mnli"
 zs_classifier = None
@@ -81,40 +156,70 @@ def load_zero_shot():
         )
 
 def classify_message(text: str) -> dict:
+    """
+    Возвращает рейтинг категорий в [-5..+5]
+    (0.5 = нейтрально, >0.5 = ближе к +5, <0.5 = ближе к -5)
+    """
     load_zero_shot()
-    # Zero‑shot принимает список меток:
     result = zs_classifier(text, candidate_labels=CATEGORIES)
-    scores = result["scores"]
     labels = result["labels"]
-    # Возвращаем словарь {label: score}
-    return {label: round(score, 3) for label, score in zip(labels, scores)}
+    scores = result["scores"]
+
+    # score 0..1 -> [-5..+5]
+    out = {}
+    for label, score in zip(labels, scores):
+        rating = (float(score) - 0.5) * 10
+        out[label] = round(rating, 2)
+    return out
 
 # =====================
 # Gradio UI
 # =====================
-with gr.Blocks(title="Unified NLP API") as demo:
-    gr.Markdown("## 📈 Unified NLP API")
-    
+with gr.Blocks(title="Unified NLP API (-5..+5)") as demo:
+    gr.Markdown("## 📈 Unified NLP API (all scores: -5 .. +5)")
+    gr.Markdown(
+        """
+**Что есть что:**
+- **Agreement**: -5 = сильное противоречие, +5 = сильное согласие  
+- **Sentiment**: -5 = негатив, +5 = позитив  
+- **Sarcasm**: -5 = уверенно *не сарказм*, +5 = уверенно *сарказм/ирония*  
+- **Categories**: рейтинг уверенности (0.5 → 0, 1.0 → +5, 0.0 → -5)
+        """
+    )
+
     # ----- Agreement Tab -----
-    with gr.Tab("Agreement"):
+    with gr.Tab("Agreement (-5..+5)"):
         msg1 = gr.Textbox(label="Message 1")
         msg2 = gr.Textbox(label="Message 2")
+
         btn_agree = gr.Button("Check Agreement")
-        out_agree = gr.Number(label="Agreement Score")
-        btn_agree.click(fn=check_agreement, inputs=[msg1, msg2], outputs=out_agree)
-    
+        out_agree = gr.Number(label="Agreement Score (-5..+5)")
+        btn_agree.click(fn=agreement_score_minus5_plus5, inputs=[msg1, msg2], outputs=out_agree)
+
+        gr.Markdown("### Agreement with Irony adjustment")
+        btn_agree_irony = gr.Button("Check Agreement (with irony)")
+        out_agree_irony = gr.Number(label="Agreement Score (irony-aware) (-5..+5)")
+        btn_agree_irony.click(fn=agreement_with_irony, inputs=[msg1, msg2], outputs=out_agree_irony)
+
     # ----- Sentiment Tab -----
-    with gr.Tab("Sentiment"):
+    with gr.Tab("Sentiment (-5..+5)"):
         text_sent = gr.Textbox(label="Text")
         btn_sent = gr.Button("Analyze Sentiment")
-        out_sent = gr.Number(label="Sentiment Score (-5 to +5)")
+        out_sent = gr.Number(label="Sentiment Score (-5..+5)")
         btn_sent.click(fn=analyze_sentiment, inputs=text_sent, outputs=out_sent)
-    
-    # ----- Multilabel (Zero‑Shot) Classification Tab -----
-    with gr.Tab("Multilabel Classification"):
+
+    # ----- Sarcasm Tab -----
+    with gr.Tab("Sarcasm / Irony (-5..+5)"):
+        text_sarc = gr.Textbox(label="Text")
+        btn_sarc = gr.Button("Analyze Sarcasm")
+        out_sarc = gr.Number(label="Sarcasm Score (-5..+5)")
+        btn_sarc.click(fn=sarcasm_score, inputs=text_sarc, outputs=out_sarc)
+
+    # ----- Multilabel (Zero-Shot) Classification Tab -----
+    with gr.Tab("Multilabel Classification (-5..+5)"):
         text_clf = gr.Textbox(label="Text")
         btn_clf = gr.Button("Classify")
-        out_clf = gr.Label(label="Categories & Scores")
+        out_clf = gr.Label(label="Categories & Scores (-5..+5)")
         btn_clf.click(fn=classify_message, inputs=text_clf, outputs=out_clf)
 
 demo.launch()