doomsday2004 commited on
Commit
7074f5d
·
verified ·
1 Parent(s): 554e2b7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -6
app.py CHANGED
@@ -2,22 +2,22 @@ import gradio as gr
2
  from sentence_transformers import CrossEncoder
3
 
4
  # Load the reranker model
5
- model = CrossEncoder("jinaai/jina-reranker-v2-base-multilingual", trust_remote_code=True )
6
 
7
- # Function to score query-document pairs
8
  def rerank(query, documents):
9
- scores = model.predict([[query, doc] for doc in documents])
 
10
  ranked_docs = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
11
  return [{"document": doc, "score": round(score, 4)} for doc, score in ranked_docs]
12
 
13
- # Define Gradio interface
14
  iface = gr.Interface(
15
  fn=rerank,
16
- inputs=["text", gr.List(gr.Textbox())],
17
  outputs="json",
18
  title="JinaAI v2 Reranker API",
19
  description="Enter a query and a list of documents. The model will rank them based on relevance.",
20
  )
21
 
22
- # Launch Gradio app
23
  iface.launch()
 
2
  from sentence_transformers import CrossEncoder
3
 
4
  # Load the reranker model
5
+ model = CrossEncoder("jinaai/jina-reranker-v2-base-multilingual", trust_remote_code=True)
6
 
7
+ # Function to rerank documents
8
  def rerank(query, documents):
9
+ documents = documents.split("\n") # Split input into a list
10
+ scores = model.predict([[query, doc] for doc in documents if doc.strip()])
11
  ranked_docs = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
12
  return [{"document": doc, "score": round(score, 4)} for doc, score in ranked_docs]
13
 
14
+ # Gradio Interface
15
  iface = gr.Interface(
16
  fn=rerank,
17
+ inputs=["text", gr.Textbox(label="Documents (One per line)", lines=5, placeholder="Enter one document per line")],
18
  outputs="json",
19
  title="JinaAI v2 Reranker API",
20
  description="Enter a query and a list of documents. The model will rank them based on relevance.",
21
  )
22
 
 
23
  iface.launch()