Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import pinecone | |
| from sentence_transformers import SentenceTransformer | |
| import torch | |
| from splade.models.transformer_rep import Splade | |
| from transformers import AutoTokenizer | |
| from datasets import load_dataset | |
| pinecone.init( | |
| api_key='884344f6-d820-4bc8-9edf-4157373df452', | |
| environment='gcp-starter' | |
| ) | |
| index = pinecone.Index('pubmed-splade') | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| # check device being run on | |
| if device != 'cuda': | |
| print("==========\n"+ | |
| "WARNING: You are not running on GPU so this may be slow.\n"+ | |
| "\n==========") | |
| dense_model = SentenceTransformer( | |
| 'msmarco-bert-base-dot-v5', | |
| device=device | |
| ) | |
| sparse_model_id = 'naver/splade-cocondenser-ensembledistil' | |
| sparse_model = Splade(sparse_model_id, agg='max') | |
| sparse_model.to(device) # move to GPU if possible | |
| sparse_model.eval() | |
| tokenizer = AutoTokenizer.from_pretrained(sparse_model_id) | |
| data = load_dataset('Binaryy/cream_listings', split='train') | |
| df = data.to_pandas() | |
| def encode(text: str): | |
| # create dense vec | |
| dense_vec = dense_model.encode(text).tolist() | |
| # create sparse vec | |
| input_ids = tokenizer(text, return_tensors='pt') | |
| with torch.no_grad(): | |
| sparse_vec = sparse_model( | |
| d_kwargs=input_ids.to(device) | |
| )['d_rep'].squeeze() | |
| # convert to dictionary format | |
| indices = sparse_vec.nonzero().squeeze().cpu().tolist() | |
| values = sparse_vec[indices].cpu().tolist() | |
| sparse_dict = {"indices": indices, "values": values} | |
| # return vecs | |
| return dense_vec, sparse_dict | |
| def search(query): | |
| dense, sparse = encode(query) | |
| # query | |
| xc = index.query( | |
| vector=dense, | |
| sparse_vector=sparse, | |
| top_k=5, # how many results to return | |
| include_metadata=True | |
| ) | |
| match_ids = [match['id'].split('-')[0] for match in xc['matches']] | |
| # Query the existing DataFrame based on 'id' | |
| filtered_df = df[df['_id'].isin(match_ids)] | |
| attributes_to_extract = ['_id', 'postedBy.accountName', 'images', 'title', 'location', 'price'] | |
| extracted_data = filtered_df[attributes_to_extract] | |
| result_json = extracted_data.to_json(orient='records') | |
| return result_json | |
| # Create a Gradio interface | |
| iface = gr.Interface( | |
| fn=search, | |
| inputs="text", | |
| outputs="json", | |
| title="Semantic Search Prototype", | |
| description="Enter your query to perform a semantic search.", | |
| ) | |
| # Launch the Gradio interface | |
| iface.launch(share=True) |