Spaces:
Runtime error
Runtime error
File size: 4,012 Bytes
e700a38 154ac9b e700a38 1576843 e700a38 44003ea e700a38 0d1fef4 e700a38 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
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) # moves 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 vecss
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', 'title', 'location', 'features', 'description', 'images',
'videos', 'available', 'price', 'attachedDocument', 'year',
'carCondition', 'engineType', 'colour', 'model', 'noOfBed',
'noOfBathroom', 'locationISO', 'forRent', 'views', 'thoseWhoSaved',
'createdAt', 'updatedAt', '__v', 'category._id', 'category.title',
'category.slug', 'category.isAdminAllowed', 'category.createdAt',
'category.updatedAt', 'category.__v', 'postedBy.pageViews.value',
'postedBy.pageViews.users', 'postedBy.totalSaved.value',
'postedBy.totalSaved.users', 'postedBy._id', 'postedBy.firstName',
'postedBy.lastName', 'postedBy.about', 'postedBy.cover',
'postedBy.email', 'postedBy.password', 'postedBy.isAdmin',
'postedBy.savedListing', 'postedBy.isVerified',
'postedBy.verifiedProfilePicture', 'postedBy.profilePicture',
'postedBy.pronoun', 'postedBy.userType', 'postedBy.accountType',
'postedBy.subscribed', 'postedBy.noOfSubscription',
'postedBy.totalListing', 'postedBy.sellerType', 'postedBy.createdAt',
'postedBy.updatedAt', 'postedBy.__v', 'postedBy.address',
'postedBy.city', 'postedBy.country', 'postedBy.gender',
'postedBy.nationality', 'postedBy.verificationType', 'postedBy.dob',
'postedBy.locationISO', 'postedBy.state', 'postedBy.zipCode',
'postedBy.otherNames', 'postedBy.facebookUrl', 'postedBy.instagramUrl',
'postedBy.phoneNumber1', 'postedBy.phoneNumber2', 'postedBy.websiteUrl',
'postedBy.accountName', 'postedBy.accountNo', 'postedBy.bankName',
'string_features', 'complete_description']
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 search.",
)
# Launch the Gradio interface
iface.launch(share=True) |