Update app.py

app.py

@@ -3,18 +3,18 @@ import torch
 from transformers import (
     AutoTokenizer,
     AutoModelForSequenceClassification,
-    XLMRobertaForSequenceClassification,
+    pipeline,
 )
 
 # =====================
 # DEVICE
 # =====================
-DEVICE = "cpu"  # HF Spaces обычно CPU
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # =====================
 # Agreement (MNLI)
 # =====================
-MNLI_MODEL = "facebook/bart-large-mnli"  # Fixed: valid public model
+MNLI_MODEL = "facebook/bart-large-mnli"
 mnli_tokenizer = None
 mnli_model = None
 
@@ -32,7 +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
+    # Считаем: entailment - contradiction
     return round((probs[2] - probs[0]).item(), 2)
 
 # =====================
@@ -57,40 +57,37 @@ 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
+    # Приводим шкалу 1–5 к -5..+5
     return round((stars - 3) * 2.5, 2)
 
 # =====================
-# Multilabel classifier
+# Zero‑Shot Classification
 # =====================
-CLASSIFIER_MODEL = "xlm-roberta-base"
+ZS_MODEL = "facebook/bart-large-mnli"
+zs_classifier = None
+
 CATEGORIES = [
     "politique", "woke", "racism", "crime",
     "police_abuse", "corruption", "hate_speech", "activism"
 ]
-clf_tokenizer = None
-clf_model = None
 
-def load_classifier():
-    global clf_tokenizer, clf_model
-    if clf_model is None:
-        clf_tokenizer = AutoTokenizer.from_pretrained(CLASSIFIER_MODEL)
-        clf_model = XLMRobertaForSequenceClassification.from_pretrained(
-            CLASSIFIER_MODEL,
-            num_labels=len(CATEGORIES),
-            problem_type="multi_label_classification"
+def load_zero_shot():
+    global zs_classifier
+    if zs_classifier is None:
+        zs_classifier = pipeline(
+            "zero-shot-classification",
+            model=ZS_MODEL,
+            device=0 if torch.cuda.is_available() else -1
         )
-        clf_model.to(DEVICE)
-        clf_model.eval()
 
 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]
-    # Return a dict for Gradio Label output
-    return {CATEGORIES[i]: float(probs[i]) for i in range(len(CATEGORIES))}
+    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)}
 
 # =====================
 # Gradio UI
@@ -113,11 +110,11 @@ with gr.Blocks(title="Unified NLP API") as demo:
         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 -----
+    # ----- Multilabel (Zero‑Shot) Classification Tab -----
     with gr.Tab("Multilabel Classification"):
         text_clf = gr.Textbox(label="Text")
         btn_clf = gr.Button("Classify")
-        out_clf = gr.Label(label="Categories")
+        out_clf = gr.Label(label="Categories & Scores")
         btn_clf.click(fn=classify_message, inputs=text_clf, outputs=out_clf)
 
 demo.launch()