Update app.py

app.py

@@ -4,31 +4,26 @@ from langchain.embeddings import HuggingFaceEmbeddings
 import torch
 
 
-def create_miread_embed(sents, bundle):
-    tokenizer = bundle[0]
-    model = bundle[1]
-    model.cpu()
-    tokens = tokenizer(sents,
-                       max_length=512,
-                       padding=True,
-                       truncation=True,
-                       return_tensors="pt"
-                       )
-    device = torch.device('cpu')
-    tokens = tokens.to(device)
-    with torch.no_grad():
-        out = model.bert(**tokens)
-        feature = out.last_hidden_state[:, 0, :]
-    return feature.cpu()
-
-
-def get_matches(query):
-    matches = vecdb.similarity_search_with_score(query, k=60)
+def get_matches1(query):
+    matches = vecdb1.similarity_search_with_score(query, k=60)
     return matches
 
+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)
+    return matches
 
-def inference(query):
-    matches = get_matches(query)
+
+def inference(query,model=1):
+    if model==1:
+        matches = get_matches1(query)
+    elif model==2:
+        matches = get_matches2(query)
+    else:
+        matches = get_matches3(query)
     auth_counts = {}
     j_bucket = {}
     n_table = []
@@ -94,17 +89,40 @@ def inference(query):
 
     return [a_output, j_output, n_output]
 
+def inference1(query):
+    return inference(query,1)
+
+def inference2(query):
+    return inference(query,2)
+
+def inference3(query):
+    return inference(query,3)
+    
 
-model_name = "biodatlab/MIReAD-Neuro-Contrastive"
+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_embedder = HuggingFaceEmbeddings(
-    model_name=model_name,
+faiss_embedder1 = HuggingFaceEmbeddings(
+    model_name=model1_name,
+    model_kwargs=model_kwargs,
+    encode_kwargs=encode_kwargs
+)
+faiss_embedder2 = HuggingFaceEmbeddings(
+    model_name=model2_name,
+    model_kwargs=model_kwargs,
+    encode_kwargs=encode_kwargs
+)
+faiss_embedder3 = HuggingFaceEmbeddings(
+    model_name=model3_name,
     model_kwargs=model_kwargs,
     encode_kwargs=encode_kwargs
 )
 
-vecdb = FAISS.load_local("nbdt_contr", faiss_embedder)
+vecdb1 = FAISS.load_local("miread_large", faiss_embedder1)
+vecdb2 = FAISS.load_local("miread_contrastive", faiss_embedder2)
+vecdb3 = FAISS.load_local("scibert_contrastive", faiss_embedder3)
 
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -114,11 +132,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     To find a recommendation, paste a `title[SEP]abstract` or `abstract` in the text box below and click \"Find Matches\".\
     Then, you can hover to authors/abstracts/journals tab to find a suggested list.\
     The data in our current demo includes authors associated with the NBDT Journal. We will update the data monthly for an up-to-date publications.")
-    gr.Markdown("**Model on Deployment: " + model_name + "**")
 
     abst = gr.Textbox(label="Abstract", lines=10)
 
-    action_btn = gr.Button(value="Find Matches")
+    action1_btn = gr.Button(value="Find Matches with MIReAD-Neuro-Large")
+    action2_btn = gr.Button(value="Find Matches with MIReAD-Neuro-Contrastive")
+    action3_btn = gr.Button(value="Find Matches with SciBERT-Neuro-Contrastive")
 
     with gr.Tab("Authors"):
         n_output = gr.Dataframe(
@@ -147,7 +166,19 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             visible=False
         )
 
-    action_btn.click(fn=inference,
+    action_btn1.click(fn=inference1,
+                     inputs=[
+                         abst,
+                     ],
+                     outputs=[a_output, j_output, n_output],
+                     api_name="neurojane")
+    action_btn2.click(fn=inference2,
+                     inputs=[
+                         abst,
+                     ],
+                     outputs=[a_output, j_output, n_output],
+                     api_name="neurojane")
+    action_btn3.click(fn=inference3,
                      inputs=[
                          abst,
                      ],