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 | |
| return Pinecone(api_key=api_key) | |
| 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 | |
| if index_name in pc.list_indexes().names(): | |
| st.info(f"Successfully connected to index '{index_name}'") | |
| return pc.Index(index_name) | |
| 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: | |
| index = connect_to_index(pc) | |
| # Load 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) | |
| # Search the index | |
| results = index.query( | |
| vector=dense_vector, | |
| top_k=5, | |
| include_metadata=True | |
| ) | |
| st.write("### Search Results:") | |
| if results.matches: | |
| for match in results.matches: | |
| score = match.score | |
| context = match.metadata.get("context", "No context available") | |
| # Display score and context in a formatted way | |
| st.markdown(f"**Score**: `{score}`") | |
| st.markdown(f"**Context**: {context}") | |
| st.markdown("---") # Divider for each result | |
| else: | |
| st.warning("No results found for this query.") | |
| else: | |
| st.error("Please enter a query and ensure the index is initialized.") | |