-
Notifications
You must be signed in to change notification settings - Fork 0
/
st_app.py
64 lines (51 loc) · 1.68 KB
/
st_app.py
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
import os
import streamlit as st
import pinecone
from sentence_transformers import SentenceTransformer
API_KEY = st.secrets['pinecone_api_key']
@st.experimental_singleton
def init_retriever():
# initialize retriever model
return SentenceTransformer('pinecone/mpnet-retriever-squad2')
@st.experimental_singleton
def init_pinecone():
pinecone.init(
api_key=API_KEY,
environment='us-west1-gcp'
)
index = pinecone.Index('qa-index')
return index
def card(id_val, source, context):
st.markdown(f"""
<div class="card" style="margin:1rem;">
<div class="card-body">
<h5 class="card-title">{source}</h5>
<h6 class="card-subtitle mb-2 text-muted">{id_val}</h6>
<p class="card-text">{context}</p>
</div>
</div>
""", unsafe_allow_html=True)
st.markdown("""
<link href="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha384-1BmE4kWBq78iYhFldvKuhfTAU6auU8tT94WrHftjDbrCEXSU1oBoqyl2QvZ6jIW3" crossorigin="anonymous">
""", unsafe_allow_html=True)
# initialize the index and retriever components
retriever = init_retriever()
index = init_pinecone()
st.write("""
# AI Q&A
Ask me a question!
""")
query = st.text_input("Search!", "")
if query != "":
# encode the query as sentence vector
xq = retriever.encode([query]).tolist()
# get relevant contexts
xc = index.query(xq, top_k=5,
include_metadata=True)
# display each context (NEW PART)
for context in xc['results'][0]['matches']:
card(
context['id'],
context['metadata']['title'],
context['metadata']['text']
)