import time import sys import streamlit as st import string from io import StringIO import pdb import json from twc_embeddings import HFModel,SimCSEModel,SGPTModel,CausalLMModel,SGPTQnAModel import torch MAX_INPUT = 100 SEM_SIMILARITY="1" DOC_RETRIEVAL="2" CLUSTERING="3" use_case = {"1":"Finding similar phrases/sentences","2":"Retrieving semantically matching information to a query. It may not be a factual match","3":"Clustering"} from transformers import BertTokenizer, BertForMaskedLM view_count_file = "view_count.txt" def get_views(): ret_val = 0 if ("view_count" not in st.session_state): try: data = int(open(view_count_file).read().strip("\n")) except: data = 0 data += 1 ret_val = data st.session_state["view_count"] = data with open(view_count_file,"w") as fp: fp.write(str(data)) else: ret_val = st.session_state["view_count"] return "{:,}".format(ret_val) def construct_model_info_for_display(model_names): options_arr = [] markdown_str = f"

Models evaluated ({len(model_names)})
" for node in model_names: options_arr .append(node["name"]) if (node["mark"] == "True"): markdown_str += f"
 • Model: {node['name']}
    Code released by: {node['orig_author']}
    Model info: {node['sota_info']['task']}
" if ("Note" in node): markdown_str += f"
    {node['Note']}link
" markdown_str += "

" markdown_str += "
Note:
• Uploaded files are loaded into non-persistent memory for the duration of the computation. They are not cached
" limit = "{:,}".format(MAX_INPUT) markdown_str += f"
• User uploaded file has a maximum limit of {limit} sentences.
" return options_arr,markdown_str st.set_page_config(page_title='TWC - Compare popular/state-of-the-art models for tasks using sentence embeddings', page_icon="logo.jpg", layout='centered', initial_sidebar_state='auto', menu_items={ 'About': 'This app was created by taskswithcode. http://taskswithcode.com' }) col,pad = st.columns([85,15]) with col: st.image("long_form_logo_with_icon.png") @st.experimental_memo def load_model(model_name,model_names): try: ret_model = None for node in model_names: if (model_name.startswith(node["name"])): obj_class = globals()[node["class"]] ret_model = obj_class() ret_model.init_model(node["model"]) assert(ret_model is not None) except Exception as e: st.error("Unable to load model:" + model_name + " " + str(e)) pass return ret_model @st.experimental_memo def cached_compute_similarity(sentences,_model,model_name,main_index): texts,embeddings = _model.compute_embeddings(sentences,is_file=False) results = _model.output_results(None,texts,embeddings,main_index) return results def uncached_compute_similarity(sentences,_model,model_name,main_index): with st.spinner('Computing vectors for sentences'): texts,embeddings = _model.compute_embeddings(sentences,is_file=False) results = _model.output_results(None,texts,embeddings,main_index) #st.success("Similarity computation complete") return results def get_model_info(model_names,model_name): for node in model_names: if (model_name == node["name"]): return node def run_test(model_names,model_name,sentences,display_area,main_index,user_uploaded): display_area.text("Loading model:" + model_name) model_info = get_model_info(model_names,model_name) if ("Note" in model_info): fail_link = f"{model_info['Note']} [link]({model_info['alt_url']})" display_area.write(fail_link) model = load_model(model_name,model_names) display_area.text("Model " + model_name + " load complete") try: if (user_uploaded): results = uncached_compute_similarity(sentences,model,model_name,main_index) else: display_area.text("Computing vectors for sentences") results = cached_compute_similarity(sentences,model,model_name,main_index) display_area.text("Similarity computation complete") return results except Exception as e: st.error("Some error occurred during prediction" + str(e)) st.stop() return {} def display_results(orig_sentences,main_index,results,response_info,app_mode): main_sent = f"
{response_info}

" score_text = "cosine_distance" if app_mode == "similarity" else "cosine_distance/score" pivot_name = "main sentence" if app_mode == "similarity" else "query" main_sent += f"
Results sorted by {score_text}. Closest to furthest away from {pivot_name}
" pivot_name = pivot_name[0].upper() + pivot_name[1:] main_sent += f"
{pivot_name}:  {orig_sentences[main_index]}
" body_sent = [] download_data = {} first = True for key in results: if (app_mode == DOC_RETRIEVAL and first): first = False continue index = orig_sentences.index(key) + 1 body_sent.append(f"
{index}] {key}   {results[key]:.2f}
") download_data[key] = f"{results[key]:.2f}" main_sent = main_sent + "\n" + '\n'.join(body_sent) st.markdown(main_sent,unsafe_allow_html=True) st.session_state["download_ready"] = json.dumps(download_data,indent=4) def init_session(): st.session_state["download_ready"] = None st.session_state["model_name"] = "ss_test" st.session_state["main_index"] = 1 st.session_state["file_name"] = "default" def app_main(app_mode,example_files,model_name_files): init_session() with open(example_files) as fp: example_file_names = json.load(fp) with open(model_name_files) as fp: model_names = json.load(fp) curr_use_case = use_case[app_mode].split(".")[0] st.markdown("
Compare popular/state-of-the-art models for tasks using sentence embeddings
", unsafe_allow_html=True) st.markdown(f"
Use cases for sentence embeddings
   •  {use_case['1']}
   •  {use_case['2']}
   •  {use_case['3']}
This app illustrates '{curr_use_case}' use case
", unsafe_allow_html=True) st.markdown(f"
views: {get_views()}
", unsafe_allow_html=True) try: with st.form('twc_form'): uploaded_file = st.file_uploader("Step 1. Upload text file(one sentence in a line) or choose an example text file below", type=".txt") selected_file_index = st.selectbox(label=f'Example files ({len(example_file_names)})', options = list(dict.keys(example_file_names)), index=0, key = "twc_file") st.write("") options_arr,markdown_str = construct_model_info_for_display(model_names) selection_label = 'Step 2. Select Model' selected_model = st.selectbox(label=selection_label, options = options_arr, index=0, key = "twc_model") st.write("") if (app_mode == "similarity"): main_index = st.number_input('Step 3. Enter index of sentence in file to make it the main sentence',value=1,min_value = 1) else: main_index = 1 st.write("") submit_button = st.form_submit_button('Run') input_status_area = st.empty() display_area = st.empty() if submit_button: start = time.time() if uploaded_file is not None: st.session_state["file_name"] = uploaded_file.name sentences = StringIO(uploaded_file.getvalue().decode("utf-8")).read() else: st.session_state["file_name"] = example_file_names[selected_file_index]["name"] sentences = open(example_file_names[selected_file_index]["name"]).read() sentences = sentences.split("\n")[:-1] if (len(sentences) < main_index): main_index = len(sentences) st.info("Selected sentence index is larger than number of sentences in file. Truncating to " + str(main_index)) if (len(sentences) > MAX_INPUT): st.info(f"Input sentence count exceeds maximum sentence limit. First {MAX_INPUT} out of {len(sentences)} sentences chosen") sentences = sentences[:MAX_INPUT] st.session_state["model_name"] = selected_model st.session_state["main_index"] = main_index results = run_test(model_names,selected_model,sentences,display_area,main_index - 1,(uploaded_file is not None)) display_area.empty() with display_area.container(): device = 'GPU' if torch.cuda.is_available() else 'CPU' response_info = f"Computation time on {device}: {time.time() - start:.2f} secs for {len(sentences)} sentences" display_results(sentences,main_index - 1,results,response_info,app_mode) #st.json(results) st.download_button( label="Download results as json", data= st.session_state["download_ready"] if st.session_state["download_ready"] != None else "", disabled = False if st.session_state["download_ready"] != None else True, file_name= (st.session_state["model_name"] + "_" + str(st.session_state["main_index"]) + "_" + '_'.join(st.session_state["file_name"].split(".")[:-1]) + ".json").replace("/","_"), mime='text/json', key ="download" ) except Exception as e: st.error("Some error occurred during loading" + str(e)) st.stop() st.markdown(markdown_str, unsafe_allow_html=True) if __name__ == "__main__": #print("comand line input:",len(sys.argv),str(sys.argv)) #app_main(sys.argv[1],sys.argv[2],sys.argv[3]) #app_main("1","sim_app_examples.json","sim_app_models.json") app_main("2","doc_app_examples.json","doc_app_models.json")