Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st | |
| from pinecone import Pinecone | |
| from sentence_transformers import SentenceTransformer | |
| # Title of the Streamlit App | |
| st.title("Pinecone Query Search on 'pubmed-splade' Index") | |
| # Initialize Pinecone globally | |
| index = None | |
| # 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 | |
| # Function to connect to the 'pubmed-splade' index | |
| def connect_to_index(pc): | |
| index_name = 'pubmed-splade' # Hardcoded index name | |
| # Connect to the 'pubmed-splade' index | |
| if index_name in pc.list_indexes().names(): | |
| st.info(f"Successfully connected to index '{index_name}'") | |
| index = pc.Index(index_name) | |
| return index | |
| else: | |
| st.error(f"Index '{index_name}' not found!") | |
| return None | |
| # Function to encode query using sentence transformers model | |
| def encode_query(model, query_text): | |
| return model.encode(query_text).tolist() | |
| # Initialize Pinecone | |
| pc = initialize_pinecone() | |
| # If Pinecone initialized successfully, proceed with index management | |
| if pc: | |
| # Connect directly to 'pubmed-splade' index | |
| index = connect_to_index(pc) | |
| # Model for query encoding | |
| model = SentenceTransformer('msmarco-bert-base-dot-v5') | |
| # Query input | |
| query_text = st.text_input("Enter a Query to Search", "Can clinicians use the PHQ-9 to assess depression?") | |
| # 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.") | |