crate app.py

app.py

@@ -0,0 +1,28 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+
+tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-v2-m3")
+model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-v2-m3")
+
+def rerank(query, docs):
+    docs = docs.strip().split('\n')
+    pairs = [(query, doc) for doc in docs]
+    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors="pt")
+    with torch.no_grad():
+        scores = model(**inputs).logits.squeeze(-1)
+    results = sorted(zip(docs, scores.tolist()), key=lambda x: x[1], reverse=True)
+    return "\n\n".join([f"Score: {score:.4f}\n{doc}" for doc, score in results])
+
+iface = gr.Interface(
+    fn=rerank,
+    inputs=[
+        gr.Textbox(label="Query", lines=1),
+        gr.Textbox(label="Documents (one per line)", lines=10)
+    ],
+    outputs="text",
+    title="BGE Reranker v2 M3",
+    description="Input a query and a list of documents. Outputs reranked documents with scores."
+)
+
+iface.launch()