Update app.py

app.py

@@ -3,32 +3,20 @@ from langchain.vectorstores import FAISS
 from langchain.embeddings import HuggingFaceEmbeddings
 import torch
 
+def get_matches1(query):
+    matches = vecdb1.similarity_search_with_score(query, k=60)
+    return matches
 
-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_matches2(query):
+    matches = vecdb2.similarity_search_with_score(query, k=60)
     return matches
 
 
-def inference(query):
-    matches = get_matches(query)
+def inference(query,method=1):
+    if method==1:
+        matches = get_matches1(query)
+    else:
+        matches = get_matches2(query)
     auth_counts = {}
     j_bucket = {}
     n_table = []
@@ -94,17 +82,29 @@ def inference(query):
 
     return [a_output, j_output, n_output]
 
+def inference1(query):
+    return inference(query,1)
+
+def inference2(query):
+    return inference(query,2)
 
-model_name = "biodatlab/MIReAD-Neuro-Large"
+model1_name = "biodatlab/MIReAD-Neuro-Large"
+model2_name = "biodatlab/MIReAD-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
 )
 
-vecdb = FAISS.load_local("nbdt_index", faiss_embedder)
+vecdb1 = FAISS.load_local("nbdt_index", faiss_embedder1)
+vecdb2 = FAISS.load_local("indexes", faiss_embedder2)
 
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -117,7 +117,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     abst = gr.Textbox(label="Abstract", lines=10)
 
-    action_btn = gr.Button(value="Find Matches")
+    action_btn = gr.Button(value="Find Matches with Normal Model")
+    action2_btn = gr.Button(value="Find Matches with Contrastive Model")
 
     with gr.Tab("Authors"):
         n_output = gr.Dataframe(
@@ -146,7 +147,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             visible=False
         )
 
-    action_btn.click(fn=inference,
+    action_btn.click(fn=inference1,
+                     inputs=[
+                         abst,
+                     ],
+                     outputs=[a_output, j_output, n_output],
+                     api_name="neurojane")
+    action2_btn.click(fn=inference2,
                      inputs=[
                          abst,
                      ],