AlekseyV commited on
Commit
b8cb8ea
·
1 Parent(s): 06ccfd7

fix output format

Browse files
Files changed (1) hide show
  1. app.py +5 -5
app.py CHANGED
@@ -6,15 +6,15 @@ ce = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")
6
  def rerank(query, docs_list):
7
  pairs = [[query, doc] for doc in docs_list]
8
  scores = ce.predict(pairs)
9
- sorted_docs = [doc for _, doc in sorted(zip(scores, docs_list), reverse=True)]
10
- return sorted_docs, scores.tolist()
11
 
12
- app = gr.Interface(
13
  fn=rerank,
14
  inputs=[gr.Textbox(label="Query"), gr.Textbox(label="Docs (JSON‑list)")],
15
- outputs=[gr.Dataframe(type="array", headers=["doc","score"])],
16
  api_name="rerank"
17
  )
18
 
19
  if __name__ == "__main__":
20
- app.launch()
 
6
  def rerank(query, docs_list):
7
  pairs = [[query, doc] for doc in docs_list]
8
  scores = ce.predict(pairs)
9
+ rows = [[doc, float(score)] for doc, score in zip(docs_list, scores)]
10
+ return rows
11
 
12
+ iface = gr.Interface(
13
  fn=rerank,
14
  inputs=[gr.Textbox(label="Query"), gr.Textbox(label="Docs (JSON‑list)")],
15
+ outputs=gr.Dataframe(type="array", headers=["doc", "score"]),
16
  api_name="rerank"
17
  )
18
 
19
  if __name__ == "__main__":
20
+ iface.launch()