Spaces:
Runtime error
Runtime error
File size: 2,295 Bytes
6779ad6 058dbba 6779ad6 058dbba 6779ad6 058dbba 6779ad6 |
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 |
import os
import logging
from haystack.pipelines.standard_pipelines import TextIndexingPipeline
from haystack.document_stores import InMemoryDocumentStore
from haystack.nodes import BM25Retriever
from haystack.nodes import FARMReader
from haystack.pipelines import ExtractiveQAPipeline
from haystack.utils import print_answers
class Question_and_Answer_System:
def __init__(self):
self.pipe_line = None
doc_dir = "data"
document_store = self.prepare_documents(doc_dir)
self.reader = self.create_reader(document_store)
self.retriever = self.create_retriever(document_store)
self.pipe_line = self.create_pipeline(self.reader, self.retriever)
def setup_logging(self):
logging.basicConfig(
format="%(levelname)s - %(name)s - %(message)s", level=logging.WARNING
)
logging.getLogger("haystack").setLevel(logging.INFO)
def prepare_documents(self, doc_dir):
document_store = InMemoryDocumentStore(use_bm25=True)
doc_dir = "data"
files_to_index = [doc_dir + "/" + f for f in os.listdir(doc_dir)]
indexing_pipeline = TextIndexingPipeline(document_store)
indexing_pipeline.run_batch(file_paths=files_to_index)
return document_store
def create_retriever(self, document_store):
retriever = BM25Retriever(document_store=document_store)
return retriever
def create_reader(self, document_store):
"""
reader = FARMReader(
model_name_or_path="deepset/roberta-base-squad2", use_gpu=True
)
"""
reader = FARMReader(
model_name_or_path="deepset/tinyroberta-squad2", use_gpu=True
)
return reader
def create_pipeline(self, reader, retriever):
self.pipe_line = ExtractiveQAPipeline(reader, retriever)
return self.pipe_line
def answer_question(self, question: str):
prediction = self.pipe_line.run(
query=question,
params={"Retriever": {"top_k": 10}, "Reader": {"top_k": 5}},
)
return prediction
def format_answers(self, prediction):
print_answers(
prediction, details="minimum" ## Choose from `minimum`, `medium`, and `all`
)
|