Update app.py

app.py

@@ -1,9 +1,12 @@
-from typing import List, Tuple
+from typing import List
+
+import numpy as np
 import pandas as pd
-from sentence_transformers import SentenceTransformer, util
 import streamlit as st
+from sentence_transformers import SentenceTransformer, util
 from st_aggrid import AgGrid, GridOptionsBuilder, JsCode
-import numpy as np
+from textblob import Sentence
+
 st.set_page_config(layout='wide')
 
 @st.cache(allow_output_mutation=True)
@@ -11,9 +14,13 @@ def load_model():
     """Load pretrained model from SentenceTransformer"""
     return SentenceTransformer('minilm_sbert')
 
-def semantic_search(model, sentence, corpus_embeddings):
+def semantic_search(model: SentenceTransformer,
+                    query: str,
+                    corpus_embeddings: List) -> pd.DataFrame:
     """Perform semantic search on the corpus"""
-    query_embeddings = model.encode(sentence,
+    query_embeddings = model.encode(sentences=query,
+                                    batch_size=128,
+                                    show_progress_bar=False,
                                     convert_to_tensor=True,
                                     normalize_embeddings=True)
 
@@ -24,24 +31,31 @@ def semantic_search(model, sentence, corpus_embeddings):
 
     return pd.DataFrame(hits[0])
 
-def get_similarity_score(model, data, query, corpus_embeddings):
-    """Get similarity score for each data point and sort by similarity score and day"""
-    hits = semantic_search(model, [query], corpus_embeddings)
+def get_similarity_score(model: Sentence,
+                         data: pd.DataFrame,
+                         query: str,
+                         corpus_embeddings: List) -> pd.DataFrame:
+    """Get similarity score for each data point and sort by similarity score and last day"""
+    hits = semantic_search(model, query, corpus_embeddings)
     result = pd.merge(data, hits, left_on='ID', right_on='corpus_id')
     result['Last Day'] = pd.to_datetime(result['Last Day'], format='%d/%m/%Y').dt.date
     result.sort_values(by=['score', 'Last Day'], ascending=[False, True], inplace=True)
     return result
 
-@st.cache(ttl=4*3600)
-def create_embedding(model: SentenceTransformer, data: pd.DataFrame, key: str) -> Tuple[list, list]:
-    """Create vector embeddings from the dataset"""
+@st.cache(ttl=2*3600)
+def create_embedding(model: SentenceTransformer,
+                     data: pd.DataFrame,
+                     key: str) -> List:
+    "Maps job title from the corpus to a 384 dimensional vector embeddings"
     corpus_sentences = data[key].astype(str).tolist()
     corpus_embeddings = model.encode(sentences=corpus_sentences,
+                                     batch_size=128,
+                                     show_progress_bar=False,
                                      convert_to_tensor=True,
                                      normalize_embeddings=True)
     return corpus_embeddings
 
-def load_dataset(columns: List) -> pd.DataFrame:
+def load_dataset(columns: List[str]) -> pd.DataFrame:
     """Load real-time dataset from google sheets"""
     sheet_id = '1KeuPPVw9gueNmMrQXk1uGFlY9H1vvhErMLiX_ZVRv_Y'
     sheet_name = 'Form Response 3'.replace(' ', '%20')
@@ -106,8 +120,6 @@ def main():
     columns = ['Timestamp', 'Full Name', 'Company', 'Previous Role',
                'Experience (months)', 'Last Day', 'LinkedIn Profile']
     data = load_dataset(columns)
-
-    # Inference
     model = load_model()
     corpus_embeddings = create_embedding(model, data, 'Previous Role')
 
@@ -115,6 +127,7 @@ def main():
     submitted = st.button('Submit')
 
     if submitted:
+        print(job_title + ',' + str(pd.Timestamp.now()))
         st.info(f'Showing results for {job_title}')
         result = get_similarity_score(model, data, job_title, corpus_embeddings)
         result = result[columns]