Spaces:
Sleeping
Sleeping
import json | |
import streamlit as st | |
from haystack.document_stores import FAISSDocumentStore | |
from haystack.document_stores import InMemoryDocumentStore | |
from transformers import DPRContextEncoder, DPRContextEncoderTokenizer | |
from haystack.nodes import DensePassageRetriever | |
from haystack.nodes import FARMReader | |
from haystack.pipelines import ExtractiveQAPipeline | |
st.header("DPR on Supreme Court Judgements (Capital Gain)") | |
with open("responses.json", 'r') as f: | |
data = json.load(f) | |
documents = [ | |
{ | |
"content": doc["text"], | |
"meta": { | |
"name": doc["title"], | |
"url": doc["url"] | |
} | |
} for doc in data | |
] | |
document_store = FAISSDocumentStore(embedding_dim=768, faiss_index_factory_str="Flat") | |
# document_store = InMemoryDocumentStore() | |
document_store.write_documents(documents) | |
retriever = DensePassageRetriever( | |
document_store=document_store, | |
query_embedding_model="facebook/dpr-question_encoder-single-nq-base", | |
passage_embedding_model="facebook/dpr-ctx_encoder-single-nq-base", | |
) | |
document_store.update_embeddings(retriever) | |
document_store.save("faiss_index") | |
reader = FARMReader(model_name_or_path="deepset/bert-base-cased-squad2") | |
pipeline = ExtractiveQAPipeline(reader=reader, retriever=retriever) | |
query = st.input() | |
# query = "What is the subject matter of the petition in the Sadanand S. Varde case?" | |
result = pipeline.run(query=query, params={"Retriever": {"top_k": 10}, "Reader": {"top_k": 5}}) | |
for answer in result['answers']: | |
st.markdown(f"=====================\nAnswer: {answer.answer}\nContext: {answer.context}\nScore: {answer.score}") |