Update app.py

app.py

@@ -1,86 +1,86 @@
 import os
 import streamlit as st
-import torch
+from pinecone import Pinecone, ServerlessSpec
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer
-from splade.models.transformer_rep import Splade
-import pinecone
 
-# 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'
+# Title of the Streamlit App
+st.title("Pinecone Index Management with Streamlit")
 
-# 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 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 using the new class instance method
+        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
 
-# Initialize Dense and Sparse models
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+# 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
 
-# Dense model (Sentence-BERT)
-dense_model = SentenceTransformer('msmarco-bert-base-dot-v5', device=device)
+# Function to encode query using sentence transformers model
+def encode_query(model, query_text):
+    return model.encode(query_text).tolist()
 
-# Sparse model (SPLADE)
-sparse_model_id = 'naver/splade-cocondenser-ensembledistil'
-sparse_model = Splade(sparse_model_id, agg='max').to(device)
-sparse_model.eval()
+# Initialize Pinecone
+pc = initialize_pinecone()
 
-# Tokenizer for sparse model
-tokenizer = AutoTokenizer.from_pretrained(sparse_model_id)
-
-# 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
-
-# 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")
+# 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)
     
-    hsparse = {
-        'indices': sparse['indices'],
-        'values':  [v * (1 - alpha) for v in sparse['values']]
-    }
-    hdense = [v * alpha for v in dense]
-    
-    return hdense, hsparse
-
-# Streamlit UI
-st.title("PubMed Search Application")
-query = st.text_input("Enter your query:", "")
+    # 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}'")
 
-# Slider to control sparse-dense scaling
-alpha = st.slider("Hybrid Search Weight (Dense vs Sparse)", 0.0, 1.0, 0.5)
+    # Model for query encoding
+    model = SentenceTransformer('msmarco-bert-base-dot-v5')
 
-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 input
+    query_text = st.text_input("Enter a Query to Search", "Can clinicians use the PHQ-9 to assess depression?")
     
-    # 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("---")
+    # 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.")
+
+# 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.")