gabor1 commited on
Commit
73dbadf
·
verified ·
1 Parent(s): d76ef2f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -10
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
- def rerank(query, docs):
9
- docs = docs.strip().split('\n')
10
- pairs = [(query, doc) for doc in docs]
 
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(docs, scores.tolist()), key=lambda x: x[1], reverse=True)
15
- return "\n\n".join([f"Score: {score:.4f}\n{doc}" for doc, score in results])
 
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="text",
24
  title="BGE Reranker v2 M3",
25
- description="Input a query and a list of documents. Outputs reranked documents with scores."
26
  )
27
 
28
- iface.launch(share=True)
 
 
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()