Create app.py

app.py

@@ -0,0 +1,79 @@
+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)