Update app.py

app.py

@@ -1,86 +1,86 @@
 import os
 import streamlit as st
-from pinecone import Pinecone, ServerlessSpec
+import torch
 from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer
+from splade.models.transformer_rep import Splade
+import pinecone
 
-# Title of the Streamlit App
-st.title("Pinecone Index Management with Streamlit")
+# Initialize Pinecone connection
+api_key = os.getenv('PINECONE_API_KEY', 'b250d1e1-fa69-40f7-81e7-442d53f62859')
+pinecone.init(api_key=api_key, environment='us-east1-gcp')
+index_name = 'pubmed-splade'
 
-# Function to initialize Pinecone
-def initialize_pinecone():
-    api_key = os.getenv('PINECONE_API_KEY')  # Get Pinecone API key from environment variable
-    if api_key:
-        # Initialize Pinecone client
-        pc = Pinecone(api_key=api_key)
-        return pc
-    else:
-        st.error("Pinecone API key not found! Please set the PINECONE_API_KEY environment variable.")
-        return None
+# Connect to the Pinecone index
+if pinecone.list_indexes() and index_name in pinecone.list_indexes():
+    index = pinecone.Index(index_name)
+else:
+    st.error("Pinecone index not found! Ensure the correct Pinecone index is being used.")
 
-# Function to create or connect to an index
-def create_or_connect_index(pc, index_name, dimension):
-    if index_name not in pc.list_indexes().names():
-        # Create index if it doesn't exist
-        pc.create_index(
-            name=index_name,
-            dimension=dimension,
-            metric='dotproduct',  # Change this based on your use case
-            spec=ServerlessSpec(cloud='aws', region='us-west-2')  # Change to your cloud provider and region
-        )
-        st.success(f"Created new index '{index_name}'")
-    else:
-        st.info(f"Index '{index_name}' already exists.")
-    # Connect to the index
-    index = pc.Index(index_name)
-    return index
+# Initialize Dense and Sparse models
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-# Function to encode query using sentence transformers model
-def encode_query(model, query_text):
-    return model.encode(query_text).tolist()
+# Dense model (Sentence-BERT)
+dense_model = SentenceTransformer('msmarco-bert-base-dot-v5', device=device)
 
-# Initialize Pinecone
-pc = initialize_pinecone()
+# Sparse model (SPLADE)
+sparse_model_id = 'naver/splade-cocondenser-ensembledistil'
+sparse_model = Splade(sparse_model_id, agg='max').to(device)
+sparse_model.eval()
 
-# If Pinecone initialized successfully, proceed with index management
-if pc:
-    index_name = st.text_input("Enter Index Name", "my_index")
-    dimension = st.number_input("Enter Vector Dimension", min_value=1, value=768)
-    
-    # Button to create or connect to index
-    if st.button("Create or Connect to Index"):
-        index = create_or_connect_index(pc, index_name, dimension)
-        if index:
-            st.success(f"Successfully connected to index '{index_name}'")
+# Tokenizer for sparse model
+tokenizer = AutoTokenizer.from_pretrained(sparse_model_id)
 
-    # Model for query encoding
-    model = SentenceTransformer('msmarco-bert-base-dot-v5')
+# Function to encode query into dense and sparse vectors
+def encode(text: str):
+    # Dense vector
+    dense_vec = dense_model.encode(text).tolist()
+    
+    # Sparse vector
+    input_ids = tokenizer(text, return_tensors='pt')
+    with torch.no_grad():
+        sparse_vec = sparse_model(d_kwargs=input_ids.to(device))['d_rep'].squeeze()
+    
+    # Extract non-zero values and indices for sparse vector
+    indices = sparse_vec.nonzero().squeeze().cpu().tolist()
+    values = sparse_vec[indices].cpu().tolist()
+    
+    sparse_dict = {"indices": indices, "values": values}
+    return dense_vec, sparse_dict
 
-    # Query input
-    query_text = st.text_input("Enter a Query to Search", "Can clinicians use the PHQ-9 to assess depression?")
+# Function for hybrid search scaling
+def hybrid_scale(dense, sparse, alpha: float):
+    if alpha < 0 or alpha > 1:
+        raise ValueError("Alpha must be between 0 and 1")
+    
+    hsparse = {
+        'indices': sparse['indices'],
+        'values':  [v * (1 - alpha) for v in sparse['values']]
+    }
+    hdense = [v * alpha for v in dense]
     
-    # Button to encode query and search the Pinecone index
-    if st.button("Search Query"):
-        if query_text and index:
-            dense_vector = encode_query(model, query_text)
-            st.write(f"Encoded Query Vector: {dense_vector}")
-            
-            # Search the index (sparse values can be added here as well)
-            results = index.query(
-                vector=dense_vector,
-                top_k=5,
-                include_metadata=True
-            )
-            
-            st.write("Search Results:")
-            for match in results.matches:
-                st.write(f"ID: {match.id}, Score: {match.score}, Metadata: {match.metadata}")
-        else:
-            st.error("Please enter a query and ensure the index is initialized.")
+    return hdense, hsparse
 
-# Option to delete index
-if st.button("Delete Index"):
-    if pc and index_name in pc.list_indexes().names():
-        pc.delete_index(index_name)
-        st.success(f"Index '{index_name}' deleted successfully.")
-    else:
-        st.error("Index not found.")
+# Streamlit UI
+st.title("PubMed Search Application")
+query = st.text_input("Enter your query:", "")
+
+# Slider to control sparse-dense scaling
+alpha = st.slider("Hybrid Search Weight (Dense vs Sparse)", 0.0, 1.0, 0.5)
+
+if query:
+    # Encode the query
+    dense_vec, sparse_vec = encode(query)
+    
+    # Scale vectors based on slider value
+    hdense, hsparse = hybrid_scale(dense_vec, sparse_vec, alpha)
+    
+    # Query Pinecone index
+    response = index.query(vector=hdense, sparse_vector=hsparse, top_k=3, include_metadata=True)
+    
+    # Display results
+    st.write(f"Top results for query: **{query}**")
+    for match in response['matches']:
+        st.write(f"**Score**: {match['score']}")
+        st.write(f"**Context**: {match['metadata']['context']}")
+        st.write("---")