List model loading support

app.py

@@ -36,7 +36,7 @@ if menu == "Sentence Similarity":
 
     inputs = []
 
-    for i in range(n_texts):
+    for i in range(int(n_texts)):
         input = st.text_input(f'Text {i + 1}:')
 
         inputs.append(input)
@@ -45,7 +45,7 @@ if menu == "Sentence Similarity":
         results = {model: inference.text_similarity(anchor, inputs, model) for model in select_models}
         df_results = {model: results[model] for model in results}
 
-        index = inputs
+        index = [f"{idx}:{input[:min(15, len(input))]}..." for idx, input in enumerate(inputs)]
         df_total = pd.DataFrame(index=index)
         for key, value in df_results.items():
             df_total[key] = list(value['score'].values)
@@ -53,7 +53,7 @@ if menu == "Sentence Similarity":
         st.write('Here are the results for selected models:')
         st.write(df_total)
         st.write('Visualize the results of each model:')
-        st.area_chart(df_total)
+        st.line_chart(df_total)
 elif menu == "Search":
     select_models = st.multiselect("Choose models", options=list(MODELS_ID), default=list(MODELS_ID)[0])
 

backend/config.py

@@ -1,6 +1,8 @@
 MODELS_ID = dict(distilroberta = 'flax-sentence-embeddings/st-codesearch-distilroberta-base',
                  mpnet = 'flax-sentence-embeddings/all_datasets_v3_mpnet-base',
                  mpnet_qa = 'flax-sentence-embeddings/mpnet_stackexchange_v1',
+                 mpnet_asymmetric_qa = ['flax-sentence-embeddings/multi-QA_v1-mpnet-asymmetric-Q',
+                                        'flax-sentence-embeddings/multi-QA_v1-mpnet-asymmetric-A'],
                  minilm_l6 = 'flax-sentence-embeddings/all_datasets_v3_MiniLM-L6')
 
 QA_MODELS_ID = dict(

backend/inference.py

@@ -1,7 +1,7 @@
 import pandas as pd
 import jax.numpy as jnp
 
-from typing import List
+from typing import List, Union
 
 # Defining cosine similarity using flax.
 from backend.utils import load_model
@@ -13,12 +13,17 @@ def cos_sim(a, b):
 
 # We get similarity between embeddings.
 def text_similarity(anchor: str, inputs: List[str], model_name: str):
+    print(model_name)
     model = load_model(model_name)
-    assert hasattr(model, 'encode')     # multiple models is not supported for similarity
 
     # Creating embeddings
-    anchor_emb = model.encode(anchor)[None, :]
-    inputs_emb = model.encode([input for input in inputs])
+    if hasattr(model, 'encode'):
+        anchor_emb = model.encode(anchor)[None, :]
+        inputs_emb = model.encode([input for input in inputs])
+    else:
+        assert len(model) == 2
+        anchor_emb = model[0].encode(anchor)[None, :]
+        inputs_emb = model[1].encode([input for input in inputs])
 
     # Obtaining similarity
     similarity = list(jnp.squeeze(cos_sim(anchor_emb, inputs_emb)))

backend/utils.py

@@ -7,10 +7,10 @@ from .config import MODELS_ID
 def load_model(model_name):
     assert model_name in MODELS_ID.keys()
     # Lazy downloading
-    models = MODELS_ID[model_name]
-    if models is str:
-        output = SentenceTransformer(models)
-    elif hasattr(models, '__iter__') :
-        output = [SentenceTransformer(model) for model in models]
+    model_ids = MODELS_ID[model_name]
+    if type(model_ids) == str:
+        output = SentenceTransformer(model_ids)
+    elif hasattr(model_ids, '__iter__'):
+        output = [SentenceTransformer(name) for name in model_ids]
 
     return output