Update app.py

app.py

@@ -11,12 +11,17 @@ def get_matches2(query):
     matches = vecdb2.similarity_search_with_score(query, k=60)
     return matches
 
+def get_matches3(query):
+    matches = vecdb3.similarity_search_with_score(query, k=60)
+
 
 def inference(query,method=1):
     if method==1:
         matches = get_matches1(query)
-    else:
+    elif method==2:
         matches = get_matches2(query)
+    else:
+        matches = get_matches3(query)
     auth_counts = {}
     j_bucket = {}
     n_table = []
@@ -88,8 +93,12 @@ def inference1(query):
 def inference2(query):
     return inference(query,2)
 
+def inference3(query):
+    return inference(query,3)
+
 model1_name = "biodatlab/MIReAD-Neuro-Large"
 model2_name = "biodatlab/MIReAD-Neuro-Contrastive"
+model3_name = "biodatlab/SciBERT-Neuro-Contrastive"
 model_kwargs = {'device': 'cpu'}
 encode_kwargs = {'normalize_embeddings': False}
 faiss_embedder1 = HuggingFaceEmbeddings(
@@ -102,9 +111,15 @@ faiss_embedder2 = HuggingFaceEmbeddings(
     model_kwargs=model_kwargs,
     encode_kwargs=encode_kwargs
 )
+faiss_embedder3 = HuggingFaceEmbeddings(
+    model_name=model3_name,
+    model_kwargs=model_kwargs,
+    encode_kwargs=encode_kwargs
+)
 
 vecdb1 = FAISS.load_local("nbdt_index", faiss_embedder1)
 vecdb2 = FAISS.load_local("indexes", faiss_embedder2)
+vecdb3 = FAISS.load_local("indexes/scibert_contr",faiss_embedder3)
 
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -118,7 +133,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     abst = gr.Textbox(label="Abstract", lines=10)
 
     action_btn = gr.Button(value="Find Matches with Normal Model")
-    action2_btn = gr.Button(value="Find Matches with Contrastive Model")
+    action2_btn = gr.Button(value="Find Matches with MIReAD Contrastive Model")
+    action3_btn = gr.Button(value="Find Matches with SciBERT Contrastive Model")
 
     with gr.Tab("Authors"):
         n_output = gr.Dataframe(
@@ -159,5 +175,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                      ],
                      outputs=[a_output, j_output, n_output],
                      api_name="neurojane")
+    action3_btn.click(fn=inference3,
+                     inputs=[
+                         abst,
+                     ],
+                     outputs=[a_output, j_output, n_output],
+                     api_name="neurojane")
 
 demo.launch(debug=True)