Update app.py

app.py

@@ -1,8 +1,14 @@
 import gradio as gr
 import torch
-from transformers import pipeline
+from transformers import AutoTokenizer, pipeline
 from typing import Dict
 
+# Custom models for zero-shot classification requiring trust_remote_code=True
+CUSTOM_MODELS = [
+    "mjwong/gte-multilingual-base-xnli",
+    "mjwong/gte-multilingual-base-xnli-anli"
+]
+
 # Available models for zero-shot classification
 AVAILABLE_MODELS = [
     "mjwong/multilingual-e5-large-instruct-xnli-anli",
@@ -10,7 +16,7 @@ AVAILABLE_MODELS = [
     "mjwong/multilingual-e5-large-xnli-anli",
     "mjwong/mcontriever-msmarco-xnli",
     "mjwong/mcontriever-xnli"
-]
+] + CUSTOM_MODELS
 
 def classify_text(
         model_name: str, 
@@ -38,7 +44,17 @@ def classify_text(
         # Set device: 0 if GPU available, else -1 for CPU
         device = 0 if torch.cuda.is_available() else -1
 
-        classifier = pipeline("zero-shot-classification", model=model_name, device=device)
+        if model_name in CUSTOM_MODELS:
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            classifier = pipeline("zero-shot-classification",
+                                  model=model_name,
+                                  tokenizer=tokenizer,
+                                  trust_remote_code=True
+                                  )
+
+        else:
+            classifier = pipeline("zero-shot-classification", model=model_name, device=device)
+
         labels_list = [label.strip() for label in labels.split(",")]
         result = classifier(text, candidate_labels=labels_list, multi_label=multi_label)
         return {label: score for label, score in zip(result["labels"], result["scores"])}