Update app.py

app.py

@@ -22,13 +22,14 @@ def create_miread_embed(sents, bundle):
     return feature.cpu()
 
 
-def get_matches(query, k):
-    matches = vecdb.similarity_search_with_score(query, k=k)
+def get_matches(query):
+    matches = vecdb.similarity_search_with_score(query, k=60)
     return matches
 
 
-def inference(query, k=30):
-    matches = get_matches(query, k)
+def inference(query):
+    matches = get_matches(query)
+    auth_counts = {}
     j_bucket = {}
     n_table = []
     a_table = []
@@ -38,7 +39,7 @@ def inference(query, k=30):
     def normaliser(x): return round(1 - (x-min_score)/max_score, 3)
     for i, match in enumerate(matches):
         doc = match[0]
-        score = round(normaliser(round(match[1].item(), 3)),3)
+        score = round(normaliser(round(match[1].item(), 3)), 3)
         title = doc.metadata['title']
         author = doc.metadata['authors'][0].title()
         date = doc.metadata.get('date', 'None')
@@ -46,7 +47,7 @@ def inference(query, k=30):
         submitter = doc.metadata.get('submitter', 'None')
         # journal = doc.metadata.get('journal', 'None').strip()
         journal = doc.metadata['journal']
-        if (journal == None or journal.strip() == ''):
+        if (journal is None or journal.strip() == ''):
             journal = 'None'
         else:
             journal = journal.strip()
@@ -63,7 +64,12 @@ def inference(query, k=30):
                   title,
                   link,
                   date]
-        n_table.append(record)
+        if auth_counts.get(author, 0) < 2:
+            n_table.append(record)
+            if auth_counts.get(author, 0) == 0:
+                auth_counts[author] = 1
+            else:
+                auth_counts[author] += 1
 
         # For abstracts
         record = [i+1,
@@ -78,8 +84,8 @@ def inference(query, k=30):
         a_table.append(record)
 
     del j_bucket['None']
-    j_table = sorted([[journal, round(score,3)] for journal,
-                     score in j_bucket.items()], 
+    j_table = sorted([[journal, round(score, 3)] for journal,
+                     score in j_bucket.items()],
                      key=lambda x: x[1], reverse=True)
     j_table = [[i+1, item[0], item[1]] for i, item in enumerate(j_table)]
     j_output = gr.Dataframe.update(value=j_table, visible=True)
@@ -98,7 +104,7 @@ faiss_embedder = HuggingFaceEmbeddings(
     encode_kwargs=encode_kwargs
 )
 
-vecdb = FAISS.load_local("nbdt_index", faiss_embedder)
+vecdb = FAISS.load_local("medline_index", faiss_embedder)
 
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -111,9 +117,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     abst = gr.Textbox(label="Abstract", lines=10)
 
-    k = gr.Slider(1, 100, step=1, value=50,
-                  label="Number of matches to consider")
-
     action_btn = gr.Button(value="Find Matches")
 
     with gr.Tab("Authors"):
@@ -146,7 +149,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     action_btn.click(fn=inference,
                      inputs=[
                          abst,
-                         k,
                      ],
                      outputs=[a_output, j_output, n_output],
                      api_name="neurojane")