Swap model to 'jinaai/jina-reranker-m0'

Added new classifier funktion to work with "simple" AutoModel.compute_score instead of zero-shot-pipeline

app.py

@@ -6,9 +6,16 @@ from transformers import pipeline
 logger = logging.getLogger("gradio_test_001")
 logger.setLevel(logging.INFO)
 logging.debug("Starting logging for gradio_test_001.")
+categories = [
+    "Legal", "Specification", "Facts and Figures",
+    "Publication", "Payment Scheme",
+    "Alternative Payment Systems", "Crypto Payments",
+    "Card Payments", "Banking", "Regulations", "Account Payments"
+]
+logging.debug("Categories to classify: " + repr(categories))
 
-classifier = pipeline("zero-shot-classification",
-                      model="facebook/bart-large-mnli")
+# classifier = pipeline("zero-shot-classification",
+#                       model="facebook/bart-large-mnli")
 
 # sequence_to_classify = "one day I will see the world"
 # candidate_labels = ['travel', 'cooking', 'dancing']
@@ -16,12 +23,6 @@ classifier = pipeline("zero-shot-classification",
 #        'doc_type.Scheme', 'content_type.Alt', 'content_type.Krypto',
 #        'content_type.Karte', 'content_type.Banking', 'content_type.Reg',
 #        'content_type.Konto']
-categories = [
-    "Legal", "Specification", "Facts and Figures",
-    "Publication", "Payment Scheme",
-    "Alternative Payment Systems", "Crypto Payments",
-    "Card Payments", "Banking", "Regulations", "Account Payments"
-]
 
 def transform_output(res: dict) -> list:
     return list(
@@ -53,13 +54,36 @@ def clf_text(txt: str | list[str]):
 # 'scores': [0.9938651323318481, 0.0032737774308770895, 0.002861034357920289],
 # 'sequence': 'one day I will see the world'}
 
+from transformers import AutoModel
+# comment out the flash_attention_2 line if you don't have a compatible GPU
+model = AutoModel.from_pretrained(
+    'jinaai/jina-reranker-m0',
+    torch_dtype="auto",
+    trust_remote_code=True,
+    # attn_implementation="flash_attention_2"
+)
+
+def clf_jina(txt: str | list[str]):
+    # construct sentence pairs
+    # text_pairs = [[query, doc] for doc in documents]
+    text_pairs = [[cat, txt] for cat in categories]
+    scores = model.compute_score(text_pairs, max_length=1024, doc_type="text")
+    return list(
+        sorted(
+            zip(categories, scores),
+            key=lambda tpl: tpl[1],
+            reverse=True
+        )
+    )
+
 
 def my_inference_function(name):
   return "Hello " + name + "!"
 
 gradio_interface = gradio.Interface(
   # fn = my_inference_function,
-  fn = clf_text,
+  # fn = clf_text,
+  clf_jina,
   inputs = "text",
   outputs = gradio.JSON()
 )