Spaces:
Runtime error
Runtime error
File size: 1,424 Bytes
02f6085 |
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 |
import streamlit as st
import faiss
from transformers import pipeline
from sentence_transformers import SentenceTransformer
import json
def load_index():
index = faiss.read_index('cdc_search.index')
return index
def load_data():
with open('./data.json') as f:
data = json.load(f)
return data
def load_embedder():
embedder = SentenceTransformer("distilbert-base-nli-stsb-mean-tokens")
return embedder
def load_qa_pipeline():
qa = pipeline("question-answering", model="ktrapeznikov/albert-xlarge-v2-squad-v2")
return qa
def load_questions():
with open('./questions.json') as f:
data = json.load(f)
return (q for q in data)
index = load_index()
embedder = load_embedder()
qa = load_qa_pipeline()
data = load_data()
def search(query: str, k=1):
encoded_query = embedder.encode([query])
top_k = index.search(encoded_query, k)
scores = top_k[0][0]
results = [data[_id] for _id in top_k[1][0]]
answers = []
for result in results:
answer = qa(question=query, context=result['text'])
if 'answer' in answer:
answers.append((answer['answer'], answer['score']))
return sorted(answers, key=lambda tup: tup[1], reverse=True)
questions = load_questions()
option = st.selectbox("Sample Questions", questions)
st.write('You selected: ', option)
st.markdown("\n".join([f"* {answer}" for (answer, _) in search(option)])) |