Navigation

app.py

@@ -6,7 +6,8 @@ from backend.config import MODELS_ID
 
 st.title('Demo using Flax-Sentence-Tranformers')
 
-st.sidebar.title('')
+st.sidebar.title('Tasks')
+menu = st.sidebar.radio("", options=["Sentence Similarity", "Search", "Clustering"], index=0)
 
 st.markdown('''
 
@@ -21,34 +22,40 @@ For more cool information on sentence embeddings, see the [sBert project](https:
 Please enjoy!!
 ''')
 
-select_models = st.multiselect("Choose models", options=list(MODELS_ID), default=list(MODELS_ID)[0])
+if menu == "Sentence Similarity":
+    select_models = st.multiselect("Choose models", options=list(MODELS_ID), default=list(MODELS_ID)[0])
 
-anchor = st.text_input(
-    'Please enter here the main text you want to compare:'
-)
+    anchor = st.text_input(
+        'Please enter here the main text you want to compare:'
+    )
 
-n_texts = st.number_input(
-    f'''How many texts you want to compare with: '{anchor}'?''',
-    value=2,
-    min_value=2)
+    n_texts = st.number_input(
+        f'''How many texts you want to compare with: '{anchor}'?''',
+        value=2,
+        min_value=2)
 
-inputs = []
+    inputs = []
 
-for i in range(n_texts):
-    input = st.text_input(f'Text {i + 1}:')
+    for i in range(n_texts):
+        input = st.text_input(f'Text {i + 1}:')
 
-    inputs.append(input)
+        inputs.append(input)
 
-if st.button('Tell me the similarity.'):
-    results = {model: inference.text_similarity(anchor, inputs, model) for model in select_models}
-    df_results = {model: results[model] for model in results}
+    if st.button('Tell me the 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
-    df_total = pd.DataFrame(index=index)
-    for key, value in df_results.items():
-        df_total[key] = list(value['score'].values)
+        index = inputs
+        df_total = pd.DataFrame(index=index)
+        for key, value in df_results.items():
+            df_total[key] = list(value['score'].values)
 
-    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.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)
+elif menu == "Search":
+    select_models = st.multiselect("Choose models", options=list(MODELS_ID), default=list(MODELS_ID)[0])
+
+elif menu == "Clustering":
+    select_models = st.multiselect("Choose models", options=list(MODELS_ID), default=list(MODELS_ID)[0])

backend/config.py

@@ -2,3 +2,9 @@ MODELS_ID = dict(distilroberta = 'flax-sentence-embeddings/st-codesearch-distilr
                  mpnet = 'flax-sentence-embeddings/all_datasets_v3_mpnet-base',
                  mpnet_qa = 'flax-sentence-embeddings/mpnet_stackexchange_v1',
                  minilm_l6 = 'flax-sentence-embeddings/all_datasets_v3_MiniLM-L6')
+
+QA_MODELS_ID = dict(
+    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']
+)

backend/inference.py

@@ -14,6 +14,7 @@ def cos_sim(a, b):
 # We get similarity between embeddings.
 def text_similarity(anchor: str, inputs: List[str], model_name: str):
     model = load_model(model_name)
+    assert hasattr(model, 'encode')     # multiple models is not supported for similarity
 
     # Creating embeddings
     anchor_emb = model.encode(anchor)[None, :]

backend/utils.py

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