Update app.py

app.py

@@ -1,25 +1,34 @@
 import os
+
+CPU_THREADS = 16
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+os.environ["OMP_NUM_THREADS"] = str(CPU_THREADS)
+os.environ["MKL_NUM_THREADS"] = str(CPU_THREADS)
+os.environ["OPENBLAS_NUM_THREADS"] = str(CPU_THREADS)
+os.environ["NUMEXPR_NUM_THREADS"] = str(CPU_THREADS)
+
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
+import numpy as np
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 MODEL_ID = "MoritzLaurer/deberta-v3-large-zeroshot-v2.0"
 
-os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
-torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", "2")))
+torch.set_num_threads(CPU_THREADS)
 torch.set_num_interop_threads(1)
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
 model.eval()
 
-clf = pipeline(
-    task="zero-shot-classification",
-    model=model,
-    tokenizer=tokenizer,
-    device=-1,
-    framework="pt",
-)
+label2id = {k.lower(): v for k, v in model.config.label2id.items()}
+entail_id = label2id.get("entailment", 2)
+
+def _softmax(x):
+    x = x - np.max(x)
+    e = np.exp(x)
+    return e / np.sum(e)
 
 def run_zero_shot(text, labels, hypothesis_template, multi_label, top_k):
     text = (text or "").strip()
@@ -33,22 +42,30 @@ def run_zero_shot(text, labels, hypothesis_template, multi_label, top_k):
     if not candidate_labels:
         return {"error": "Enter at least 1 label (comma-separated)."}
 
+    scores = []
     with torch.inference_mode():
-        out = clf(
-            sequences=text,
-            candidate_labels=candidate_labels,
-            hypothesis_template=hypothesis_template,
-            multi_label=bool(multi_label),
-        )
-
-    pairs = list(zip(out["labels"], out["scores"]))
+        for lab in candidate_labels:
+            hyp = hypothesis_template.format(lab)
+            inputs = tokenizer(text, hyp, return_tensors="pt", truncation=True)
+            logits = model(**inputs).logits[0].float().cpu().numpy()
+            score = float(_softmax(logits)[entail_id])
+            scores.append(score)
+
+    scores_np = np.array(scores, dtype=np.float32)
+
+    if bool(multi_label):
+        out_scores = scores_np
+    else:
+        out_scores = _softmax(scores_np)
+
+    pairs = list(zip(candidate_labels, out_scores.tolist()))
     pairs.sort(key=lambda x: x[1], reverse=True)
     pairs = pairs[: max(1, int(top_k))]
 
     return {
+        "cpu_threads": CPU_THREADS,
         "top": {"label": pairs[0][0], "confidence_pct": round(pairs[0][1] * 100, 2)},
         "all": [{"label": k, "confidence_pct": round(v * 100, 2)} for k, v in pairs],
-        "raw": out,
     }
 
 demo = gr.Interface(
@@ -61,7 +78,7 @@ demo = gr.Interface(
         gr.Slider(label="Top-K to show", minimum=1, maximum=25, value=5, step=1),
     ],
     outputs=gr.JSON(label="Output"),
-    title="Zero-Shot Classification (DeBERTa v3 Large, MoritzLaurer)",
+    title="Zero-Shot Classification (DeBERTa v3 Large, 16-core CPU)",
     flagging_mode="never",
 )