changes

app.py

@@ -1,56 +1,40 @@
 import gradio as gr
 from sentence_transformers import SentenceTransformer, util, CrossEncoder
-import torch
 from transformers import set_seed
 import numpy as np
-import pandas as pd
-import argparse
 
 set_seed(42)
 
+   
 
-def calc_preferred_dense(doc1, doc2, q1, q2, model_name="dpr", model=None):
+passage_encoder = SentenceTransformer(
+    "facebook-dpr-ctx_encoder-multiset-base"
+)
+query_encoder = SentenceTransformer(
+    "facebook-dpr-question_encoder-multiset-base"
+)
+model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")
+embedder = SentenceTransformer("all-mpnet-base-v2")
+
+
+def calc_preferred_dense(doc1, doc2, q1, q2, model_name="dpr"):
     """
     Input:
         doc1, doc2: strings containing the documents/passages
         query1, query2: strings for queries that are only relevant to the corresponding doc (doc1 -> q1, doc2 -> q2)
         model_name: string containing the type of model to run
-        model: the preloaded model, if caching
 
     Returns:
         A dictionary containing each query (q1 or q2) and the score (P@1) for the pair
 
     """
-    ### Model initialization
-    if model_name == "dpr":
-        model_type = "dpr"
-        if model is not None:
-            passage_encoder, query_encoder = model
-        else:
-            passage_encoder = SentenceTransformer(
-                "facebook-dpr-ctx_encoder-multiset-base"
-            )
-            query_encoder = SentenceTransformer(
-                "facebook-dpr-question_encoder-multiset-base"
-            )
-    elif "cross-encoder" in model_name or "t5" in model_name:
-        model_type = "cross_encoder"
-        if model is None:
-            model = CrossEncoder(model_name)
-    else:
-        model_type = "biencoder"
-        if model is not None:
-            embedder = model
-        else:
-            embedder = SentenceTransformer(model_name)
-
     corpus = [doc1, doc2]
     queries = [q1, q2]
     results = {}
     num_correct = 0
 
     ### Do Retrieval
-    if model_type == "dpr":
+    if model_name == "dpr":
         passage_embeddings = passage_encoder.encode(corpus)
 
         query_encoder = SentenceTransformer(
@@ -69,7 +53,7 @@ def calc_preferred_dense(doc1, doc2, q1, q2, model_name="dpr", model=None):
                 num_correct += 1
         model = (passage_encoder, query_encoder)
 
-    elif model_type == "cross_encoder":
+    elif model_name == "cross_encoder":
         for idx, query in enumerate(queries):
             scores = model.predict([[query, doc1], [query, doc2]])
             results[f"q{idx+1}"] = scores.tolist()
@@ -100,11 +84,18 @@ def calc_preferred_dense(doc1, doc2, q1, q2, model_name="dpr", model=None):
         model = embedder
 
     results["score"] = num_correct / 2
-    return results, model
+    return results
 
 
-def greet(name):
-    return "Hello " + name + "!!"
+gr.Interface(
+    calc_preferred_dense,
+    [ gr.Textbox(label="Sentence 1"), gr.Textbox(label="Sentence 2"), gr.Dropdown(["dpr", "cross-encoder", "dense"], value="cross-encoder")],
+    [ gr.components.Label(label="Similarity score") ],
+    title="Similarity score between 2 sentences",
+    description="In this demo do provide 2 sentences bellow. They can even be in distinct languages. Powered by S-BERT multilingual model : https://www.sbert.net.",
+    examples=[['The sentences are mapped such that sentences with similar meanings are close in vector space.', 'Les phrases sont mappées de manière à ce que les phrases ayant des significations similaires soient proches dans l\'espace vectoriel.'],
+              ['You do not need to specify the input language.', 'You can use any language.']],
+    live=True,
+    allow_flagging="never"
+).launch(debug=True, enable_queue=True)
 
-iface = gr.Interface(fn=calc_preferred_dense, inputs="text", outputs="text")
-iface.launch()

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+sentence_transformers
+numpy