Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import pandas as pd | |
| from safetensors import safe_open | |
| from sentence_transformers import SentenceTransformer, CrossEncoder, util | |
| import pickle | |
| st.title('Search offers') | |
| bi_encoder = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') | |
| cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') | |
| tensors = {} | |
| with safe_open("embeddings.safetensors", framework="pt") as f : | |
| for k in f.keys(): | |
| tensors[k] = f.get_tensor(k) | |
| corpus_embeddings = tensors['embedding'] | |
| with open('corpus.pickle', 'rb') as f: | |
| passages = pickle.load(f) | |
| def search(query, top_k): | |
| query_embedding = bi_encoder.encode(query, convert_to_tensor=True) | |
| query_embedding = query_embedding #.cuda() | |
| hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=top_k) | |
| hits = hits[0] | |
| cross_inp = [[query, passages[hit['corpus_id']]] for hit in hits] | |
| cross_scores = cross_encoder.predict(cross_inp) | |
| # Sort results by the cross-encoder scores | |
| for idx in range(len(cross_scores)): | |
| hits[idx]['cross-score'] = cross_scores[idx] | |
| hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True) | |
| score_list, output_list = [],[] | |
| for hit in hits[:10]: | |
| score_list.append("{:.3f}".format(hit['cross-score'])) | |
| temp_output = passages[hit['corpus_id']].replace("\n", " ") | |
| temp_output = list(temp_output.rsplit('{'))[0].strip() | |
| output_list.append(temp_output) | |
| dataframe = pd.DataFrame({'score': score_list, 'offers': output_list}) | |
| dataframe.drop_duplicates(subset=['offers'], keep='first', inplace=True) | |
| return dataframe | |
| with st.form("my_form"): | |
| query = st.text_input("Enter the brand name, category or retailer name to search \ | |
| for relevant offers 👇", | |
| placeholder = "Enter the text here") | |
| num = st.number_input('Manximum number of offers to display', min_value=1, max_value=10) | |
| submitted = st.form_submit_button("Submit") | |
| if submitted: | |
| df = search(query, num) | |
| st.dataframe(df, use_container_width=True) |