Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -2,27 +2,32 @@ import gradio as gr
|
|
2 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
3 |
import torch
|
4 |
|
|
|
5 |
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-v2-m3")
|
6 |
model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-v2-m3")
|
7 |
|
8 |
-
|
9 |
-
|
10 |
-
|
|
|
11 |
inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors="pt")
|
12 |
with torch.no_grad():
|
13 |
scores = model(**inputs).logits.squeeze(-1)
|
14 |
-
results = sorted(zip(
|
15 |
-
|
|
|
16 |
|
|
|
17 |
iface = gr.Interface(
|
18 |
fn=rerank,
|
19 |
inputs=[
|
20 |
-
gr.Textbox(label="Query", lines=1),
|
21 |
-
gr.Textbox(label="Documents (one per line)", lines=10)
|
22 |
],
|
23 |
-
outputs="
|
24 |
title="BGE Reranker v2 M3",
|
25 |
-
description="Input a query and
|
26 |
)
|
27 |
|
28 |
-
|
|
|
|
2 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
3 |
import torch
|
4 |
|
5 |
+
# Load tokenizer and model
|
6 |
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-v2-m3")
|
7 |
model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-v2-m3")
|
8 |
|
9 |
+
# Define reranking function
|
10 |
+
def rerank(query, documents_text):
|
11 |
+
documents = documents_text.strip().split('\n')
|
12 |
+
pairs = [(query, doc) for doc in documents]
|
13 |
inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors="pt")
|
14 |
with torch.no_grad():
|
15 |
scores = model(**inputs).logits.squeeze(-1)
|
16 |
+
results = sorted(zip(documents, scores.tolist()), key=lambda x: x[1], reverse=True)
|
17 |
+
output = "\n\n".join([f"Score: {score:.4f}\n{doc}" for doc, score in results])
|
18 |
+
return output
|
19 |
|
20 |
+
# Gradio Interface
|
21 |
iface = gr.Interface(
|
22 |
fn=rerank,
|
23 |
inputs=[
|
24 |
+
gr.Textbox(label="Query", placeholder="Enter your search query", lines=1),
|
25 |
+
gr.Textbox(label="Documents (one per line)", placeholder="Enter one document per line", lines=10)
|
26 |
],
|
27 |
+
outputs=gr.Textbox(label="Reranked Output"),
|
28 |
title="BGE Reranker v2 M3",
|
29 |
+
description="Input a query and multiple documents. Returns reranked results with scores."
|
30 |
)
|
31 |
|
32 |
+
# Launch the interface (no share=True needed)
|
33 |
+
iface.launch()
|