TheresaQWQ commited on
Commit
bd04035
·
verified ·
1 Parent(s): 0ed0161

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -3
app.py CHANGED
@@ -8,6 +8,7 @@ model = AutoModelForSequenceClassification.from_pretrained(
8
  trust_remote_code=True,
9
  )
10
 
 
11
  model.eval()
12
 
13
  def compute_scores(query, documents):
@@ -21,7 +22,8 @@ def compute_scores(query, documents):
21
  Returns:
22
  list of float: Scores representing the relevance of each document to the query.
23
  """
24
- sentence_pairs = [[query, doc] for doc in documents]
 
25
  scores = model.compute_score(sentence_pairs, max_length=1024)
26
  return scores.tolist()
27
 
@@ -29,8 +31,8 @@ def compute_scores(query, documents):
29
  iface = gr.Interface(
30
  fn=compute_scores,
31
  inputs=[
32
- gr.inputs.Textbox(lines=2, placeholder="Enter your query here..."),
33
- gr.inputs.Textbox(lines=8, placeholder="Enter your documents separated by newlines...")
34
  ],
35
  outputs="json",
36
  title="Sentence Pair Scoring with Jina Reranker Model",
 
8
  trust_remote_code=True,
9
  )
10
 
11
+ model.to('cuda') # or 'cpu' if no GPU is available
12
  model.eval()
13
 
14
  def compute_scores(query, documents):
 
22
  Returns:
23
  list of float: Scores representing the relevance of each document to the query.
24
  """
25
+ documents_list = documents.split('\n')
26
+ sentence_pairs = [[query, doc] for doc in documents_list]
27
  scores = model.compute_score(sentence_pairs, max_length=1024)
28
  return scores.tolist()
29
 
 
31
  iface = gr.Interface(
32
  fn=compute_scores,
33
  inputs=[
34
+ gr.Textbox(lines=2, placeholder="Enter your query here..."),
35
+ gr.Textbox(lines=8, placeholder="Enter your documents separated by newlines...")
36
  ],
37
  outputs="json",
38
  title="Sentence Pair Scoring with Jina Reranker Model",