Update app.py

app.py

@@ -1,23 +1,27 @@
 import gradio as gr
-from sentence_transformers import CrossEncoder
+from optimum.pipelines import pipeline
+from transformers import AutoTokenizer
 
-# Load the reranker model
-model = CrossEncoder("jinaai/jina-reranker-v2-base-multilingual", trust_remote_code=True)
+# Load ONNX optimized model
+model_name = "jinaai/jina-reranker-v2-base-multilingual"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = pipeline("text-classification", model=model_name, tokenizer=tokenizer)
 
 # Function to rerank documents
 def rerank(query, documents):
-    documents = documents.split("&&")  # Use custom delimiter to split documents correctly
-    scores = model.predict([[query, doc] for doc in documents if doc.strip()])
-    ranked_docs = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
+    documents = documents.split("&&&")
+    inputs = [[query, doc] for doc in documents if doc.strip()]
+    scores = model(inputs)
+    ranked_docs = sorted(zip(documents, [s['score'] for s in scores]), key=lambda x: x[1], reverse=True)
     return [{"document": doc, "score": round(score, 4)} for doc, score in ranked_docs]
 
 # Gradio Interface
 iface = gr.Interface(
     fn=rerank,
-    inputs=["text", gr.Textbox(label="Documents (Separate with &&)", placeholder="Doc1 ||| Doc2 ||| Doc3")], 
+    inputs=["text", gr.Textbox(label="Documents (Separate with &&&)", placeholder="Doc1 &&& Doc2 &&& Doc3")], 
     outputs="json",
-    title="JinaAI v2 Reranker API",
-    description="Enter a query and documents (separated by '&&'). The model will rank them based on relevance.",
+    title="JinaAI v2 Reranker API (Optimized)",
+    description="Enter a query and documents (separated by '&&&'). The model will rank them based on relevance.",
 )
 
 iface.launch()