Restructure inference steps

app.py

@@ -1,116 +1,103 @@
+from typing import List, Tuple
 import pandas as pd
 from sentence_transformers import SentenceTransformer, util
 import streamlit as st
 from st_aggrid import AgGrid, GridOptionsBuilder, JsCode
-from cpu_unpickler import cpu_unpickler
 st.set_page_config(layout='wide')
 
 @st.cache(allow_output_mutation=True)
 def load_model():
     return SentenceTransformer('all-MiniLM-L6-v2')
 
-def find_top_similar(sentence, corpus_sentences, corpus_embeddings):
-    # preprocess query
-    model = load_model()
-    query_embeddings = model.encode(sentence, convert_to_tensor=True) # encode to tensor
-    # query_embeddings = query_embeddings.to('cuda') # put into gpu
-    query_embeddings = util.normalize_embeddings(query_embeddings) # normalize
+def semantic_search(model, sentence, corpus_embeddings):
+    query_embeddings = model.encode(sentence,
+                                    convert_to_tensor=True,
+                                    normalize_embeddings=True)
 
-    # find the closest 5 sentences of the corpus for each query sentence based on cosine similarity
     hits = util.semantic_search(query_embeddings,
                                 corpus_embeddings,
                                 top_k=len(corpus_embeddings),
                                 score_function=util.dot_score)
-    hits = hits[0] # get the hits for the first query
-
-    # Create dataframe to store top searches
-    records = []
-
-    for hit in hits[0:len(corpus_embeddings)]:
-        records.append(corpus_sentences[hit['corpus_id']])
-
-    return records
 
-def top_k_similarity(df, query, corpus_sentences, corpus_embeddings):
-    hits = find_top_similar([query], corpus_sentences, corpus_embeddings)
+    return pd.DataFrame(hits[0])
 
-    res = pd.DataFrame()
-
-    for h in hits:
-        s = df[df['Last job role'] == h]
-        res = pd.concat([res, s])
-
-    return res
-
-def get_result(df, query, corpus_sentences, corpus_embeddings):
-    result = top_k_similarity(df, query, corpus_sentences, corpus_embeddings)
-    result.drop_duplicates(inplace=True)
+def top_k_similarity(model, df, query, corpus_embeddings):
+    hits = semantic_search(model, [query], corpus_embeddings)
+    result = pd.merge(df, hits, left_on='ID', right_on='corpus_id')
+    result.sort_values(by='score', ascending=False, inplace=True)
     return result
 
 @st.cache(allow_output_mutation=True)
-def load_embedding():
-    """Loads the embeddings from the pickle file"""
-    with open('corpus_embeddings.pkl', 'rb') as file:
-        cache_data = cpu_unpickler(file).load()
-        corpus_sentences = cache_data['sentences']
-        corpus_embeddings = cache_data['embeddings']
-
-    return corpus_sentences, corpus_embeddings
-
-def main():
-    # get dataset
+def create_embedding(model: SentenceTransformer, data: pd.DataFrame, key: str) -> Tuple[list, list]:
+    """Create vector embeddings from the dataset"""
+    corpus_sentences = data[key].astype(str).tolist()
+    corpus_embeddings = model.encode(sentences=corpus_sentences,
+                                     show_progress_bar=True,
+                                     convert_to_tensor=True,
+                                     normalize_embeddings=True)
+    return corpus_embeddings
+
+def load_dataset(columns: List) -> pd.DataFrame:
+    """Load real-time dataset from google sheets"""
     sheet_id = '1KeuPPVw9gueNmMrQXk1uGFlY9H1vvhErMLiX_ZVRv_Y'
     sheet_name = 'Form Response 3'.replace(' ', '%20')
     url = f'https://docs.google.com/spreadsheets/d/{sheet_id}/gviz/tq?tqx=out:csv&sheet={sheet_name}'
-    df = pd.read_csv(url)
-    df  = df.iloc[: , :7]
+    data = pd.read_csv(url)
+    data  = data.iloc[: , :7]
+    data.columns = columns
+    data.insert(0, 'ID', range(len(data)))
+    return data
+
+def show_aggrid_table(result: pd.DataFrame):
+    gb = GridOptionsBuilder.from_dataframe(result)
+    gb.configure_pagination(paginationAutoPageSize=True)
+    gb.configure_side_bar()
+    gb.configure_selection('multiple', use_checkbox=True, groupSelectsChildren="Group checkbox select children")
+    gb.configure_column(field="LinkedIn Profile",
+                        headerName="LinkedIn Profile",
+                        cellRenderer=JsCode('''function(params) {return `<a href=${params.value} target="_blank">${params.value}</a>`}'''))
+
+    gridOptions = gb.build()
+
+    grid_response = AgGrid(
+        dataframe=result,
+        gridOptions=gridOptions,
+        height=1100,
+        fit_columns_on_grid_load=True,
+        data_return_mode='AS_INPUT',
+        update_mode='VALUE_CHANGED',
+        theme='light',
+        enable_enterprise_modules=True,
+        allow_unsafe_jscode=True,
+    )
 
-    # get embeddings
-    corpus_sentences, corpus_embeddings = load_embedding()
-
-    # streamlit form
+def main():
     st.title('Job Posting Similarity')
+    st.write('This app will help you find similar job titles real-time from ecommurz google sheets.')
+
+    columns = ['Timestamp', 'Full Name', 'Company', 'Previous Role',
+               'Experience', 'Last Day', 'LinkedIn Profile']
+    data = load_dataset(columns)
+    model = load_model()
+    corpus_embeddings = create_embedding(model, data, 'Previous Role')
+
     job_title = st.text_input('Insert the job title below:', '')
     submitted = st.button('Submit')
 
     if submitted:
-        st.info(f'Showing results for { job_title}')
-        result = get_result(df, job_title, corpus_sentences, corpus_embeddings)
-        result.reset_index(drop=True, inplace=True)
-        result.index += 1
+        st.info(f'Showing results for {job_title}')
+        result = top_k_similarity(model, data, job_title, corpus_embeddings)
+        result = result[columns]
 
         st.download_button(
-            "Press to Download",
+            "Download Table",
             result.to_csv().encode('utf-8'),
             "result.csv",
             "text/csv",
             key='download-csv'
         )
 
-        gb = GridOptionsBuilder.from_dataframe(result)
-        gb.configure_pagination(paginationAutoPageSize=True) # Add pagination
-        # gb.configure_side_bar() #Add a sidebar
-        # gb.configure_selection('multiple', use_checkbox=True, groupSelectsChildren="Group checkbox select children") #Enable multi-row selection
-
-        gb.configure_column("LinkedIn Link",
-                            headerName="LinkedIn Link",
-                            # cellRenderer=JsCode('''function(params) {return '<a href=params.value + '" target="_blank">'+ params.value+'</a>'}'''),
-                            cellRenderer=JsCode('''function(params) {return `<a href=${params.value} target="_blank">${params.value}</a>`}'''),
-                            width=300)
-
-        gridOptions = gb.build()
-
-        grid_response = AgGrid(
-            dataframe=result,
-            gridOptions=gridOptions,
-            height=1100,
-            fit_columns_on_grid_load=True,
-            data_return_mode='AS_INPUT',
-            update_mode='VALUE_CHANGED',
-            theme='light',
-            enable_enterprise_modules=True,
-            allow_unsafe_jscode=True,
-        )
+        show_aggrid_table(result)
 
 if __name__ == '__main__':
     main()