Asymmetric QA

app.py

@@ -2,12 +2,12 @@ import streamlit as st
 import pandas as pd
 
 from backend import inference
-from backend.config import MODELS_ID
+from backend.config import MODELS_ID, QA_MODELS_ID
 
 st.title('Demo using Flax-Sentence-Tranformers')
 
 st.sidebar.title('Tasks')
-menu = st.sidebar.radio("", options=["Sentence Similarity", "Search", "Clustering"], index=0)
+menu = st.sidebar.radio("", options=["Sentence Similarity", "Asymmetric QA", "Search", "Clustering"], index=0)
 
 st.markdown('''
 
@@ -42,7 +42,7 @@ if menu == "Sentence Similarity":
         inputs.append(input)
 
     if st.button('Tell me the similarity.'):
-        results = {model: inference.text_similarity(anchor, inputs, model) for model in select_models}
+        results = {model: inference.text_similarity(anchor, inputs, model, MODELS_ID) for model in select_models}
         df_results = {model: results[model] for model in results}
 
         index = [f"{idx}:{input[:min(15, len(input))]}..." for idx, input in enumerate(inputs)]
@@ -54,6 +54,41 @@ if menu == "Sentence Similarity":
         st.write(df_total)
         st.write('Visualize the results of each model:')
         st.line_chart(df_total)
+elif menu == "Asymmetric QA":
+    select_models = st.multiselect("Choose models", options=list(QA_MODELS_ID), default=list(QA_MODELS_ID)[0])
+
+    anchor = st.text_input(
+        'Please enter here the query you want to compare with given answers:',
+        value="How many close friends do you have?"
+    )
+
+    n_texts = st.number_input(
+        f'''How many answers you want to compare with: '{anchor}'?''',
+        value=3,
+        min_value=2)
+
+    inputs = []
+
+    defaults = ["I have 10.", "How many children do you have?", "I have 3 brothers."]
+    for i in range(int(n_texts)):
+        input = st.text_input(f'Answer {i + 1}:', value=defaults[i] if i < len(defaults) else "")
+
+        inputs.append(input)
+
+    if st.button('Tell me Answer likeliness.'):
+        results = {model: inference.text_similarity(anchor, inputs, model, QA_MODELS_ID) for model in select_models}
+        df_results = {model: results[model] for model in results}
+
+        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)
+
+        st.write('Here are the results for selected models:')
+        st.write(df_total)
+        st.write('Visualize the results of each model:')
+        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,12 +1,10 @@
 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(
     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']
+                           'flax-sentence-embeddings/multi-QA_v1-mpnet-asymmetric-A'],
+    distilbert_qa = 'flax-sentence-embeddings/multi-qa_v1-distilbert-cls_dot'
 )

backend/inference.py

@@ -4,6 +4,7 @@ import jax.numpy as jnp
 from typing import List, Union
 
 # Defining cosine similarity using flax.
+from backend.config import MODELS_ID
 from backend.utils import load_model
 
 
@@ -12,9 +13,9 @@ def cos_sim(a, b):
 
 
 # We get similarity between embeddings.
-def text_similarity(anchor: str, inputs: List[str], model_name: str):
+def text_similarity(anchor: str, inputs: List[str], model_name: str, model_dict: dict):
     print(model_name)
-    model = load_model(model_name)
+    model = load_model(model_name, model_dict)
 
     # Creating embeddings
     if hasattr(model, 'encode'):

backend/utils.py

@@ -1,13 +1,12 @@
 import streamlit as st
 from sentence_transformers import SentenceTransformer
-from .config import MODELS_ID
 
 
 @st.cache(allow_output_mutation=True)
-def load_model(model_name):
-    assert model_name in MODELS_ID.keys()
+def load_model(model_name, model_dict):
+    assert model_name in model_dict.keys()
     # Lazy downloading
-    model_ids = MODELS_ID[model_name]
+    model_ids = model_dict[model_name]
     if type(model_ids) == str:
         output = SentenceTransformer(model_ids)
     elif hasattr(model_ids, '__iter__'):