Spaces:
Runtime error
Runtime error
File size: 2,406 Bytes
ee69da4 |
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 |
from haystack.nodes.retriever import EmbeddingRetriever
from haystack.nodes import TableReader, FARMReader, RouteDocuments, JoinAnswers
from haystack import Pipeline
text_reader_types = {
"minilm": "deepset/minilm-uncased-squad2",
"distilroberta": "deepset/tinyroberta-squad2",
"electra-base": "deepset/electra-base-squad2",
"bert-base": "deepset/bert-base-cased-squad2",
"deberta-large": "deepset/deberta-v3-large-squad2",
"gpt3": "implement openai answer generator"
}
table_reader_types = {
"tapas": "deepset/tapas-large-nq-hn-reader",
"text": "implement changing tables to text"
}
def create_retriever(document_store):
retriever = EmbeddingRetriever(document_store=document_store, embedding_model="deepset/all-mpnet-base-v2-table")
document_store.update_embeddings(retriever=retriever)
return document_store, retriever
def create_readers_and_pipeline(retriever, text_reader_type = "deepset/roberta-base-squad2", table_reader_type="deepset/tapas-large-nq-hn-reader", use_table=True, use_text=True):
both = (use_table and use_text)
if use_text or both:
print("Initializing Text reader..")
text_reader = FARMReader(text_reader_type)
if use_table or both:
print("Initializing table reader..")
table_reader = TableReader(table_reader_type)
if both:
route_documents = RouteDocuments()
join_answers = JoinAnswers()
text_table_qa_pipeline = Pipeline()
text_table_qa_pipeline.add_node(component=retriever, name="EmbeddingRetriever", inputs=["Query"])
if use_table and not use_text:
text_table_qa_pipeline.add_node(component=table_reader, name="TableReader", inputs=["EmbeddingRetriever"])
elif use_text and not use_table:
text_table_qa_pipeline.add_node(component=text_reader, name="TextReader", inputs=["EmbeddingRetriever"])
elif both:
text_table_qa_pipeline.add_node(component=route_documents, name="RouteDocuments", inputs=["EmbeddingRetriever"])
text_table_qa_pipeline.add_node(component=text_reader, name="TextReader", inputs=["RouteDocuments.output_1"])
text_table_qa_pipeline.add_node(component=table_reader, name="TableReader", inputs=["RouteDocuments.output_2"])
text_table_qa_pipeline.add_node(component=join_answers, name="JoinAnswers", inputs=["TextReader", "TableReader"])
return text_table_qa_pipeline
|