File size: 2,648 Bytes
e2b9039
 
3ecbeff
e2b9039
 
b21411e
a19ad68
e2b9039
d792c52
 
 
b21411e
 
 
 
 
a19ad68
b21411e
 
 
e2b9039
3ecbeff
 
 
 
a19ad68
 
b21411e
3ecbeff
 
e2b9039
b21411e
 
 
e2b9039
b21411e
fff69e7
e2b9039
b21411e
 
3ecbeff
b0a56a3
a19ad68
b21411e
b0a56a3
b21411e
 
b0a56a3
b21411e
 
 
 
a19ad68
 
b21411e
 
a19ad68
b21411e
 
 
a19ad68
 
 
 
 
 
 
 
 
 
 
 
b21411e
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
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.")