Gideon / app.py
cools's picture
Update app.py
9473952
raw
history blame contribute delete
No virus
1.52 kB
import streamlit as st
import os
import requests
import pandas as pd
import boto3
import AWSHandler
import pinecone
from sentence_transformers import SentenceTransformer
aws_access_key = os.getenv("aws_access_key")
aws_secret_key = os.getenv("aws_secret_key")
pinecone_api_key = os.getenv("pinecone_api_key")
pinecone_environment = os.getenv("pinecone_environment")
s3 = boto3.client('s3', aws_access_key_id=aws_access_key, aws_secret_access_key=aws_secret_key)
model = SentenceTransformer('all-mpnet-base-v2')
st.set_page_config(layout="wide")
def search_index(query):
pinecone.init(api_key=pinecone_api_key, environment=pinecone_environment)
index = pinecone.Index("scotus")
vector = model.encode([query]).tolist()
print(index)
print(index.query)
print(index.query(vector=vector,top_k=10,include_values=False,include_metadata=True).to_dict())
return index.query(vector=vector,top_k=10,include_values=False,include_metadata=True).to_dict()['matches']
st.markdown('''
<style>
[data-testid="stMarkdownContainer"] ul{
list-style-position: inside;
}
</style>
''', unsafe_allow_html=True)
search = st.text_input('Search', placeholder='First Amendment cannot protect libel')
if search != "":
res = search_index(query=search)
for m in res:
text = m['metadata']['Text']
if 'Name' in m['metadata']:
name = m['metadata']['Name']
else:
name = ""
st.markdown("- " + text + ' ' + '(_' + name + '_)')