mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-08-30 03:16:46 +00:00
165a0a5faa
* mock all translator tests and move one to e2e * typo * extract pipeline tests using translator * remove duplicate test * move generator test in e2e * Update e2e/pipelines/test_extractive_qa.py * pytest.mark.unit * black * remove model name as well * remove unused fixture * rename original and improve pipeline tests * fixes * pylint
30 lines
1.4 KiB
Python
30 lines
1.4 KiB
Python
from haystack.nodes import TransformersTranslator, FARMReader, TfidfRetriever
|
|
from haystack.pipelines import ExtractiveQAPipeline, TranslationWrapperPipeline
|
|
from haystack.document_stores import InMemoryDocumentStore
|
|
|
|
|
|
def test_extractive_qa_answers_with_translator(docs):
|
|
en_to_de_translator = TransformersTranslator(model_name_or_path="Helsinki-NLP/opus-mt-en-de")
|
|
de_to_en_translator = TransformersTranslator(model_name_or_path="Helsinki-NLP/opus-mt-de-en")
|
|
|
|
ds = InMemoryDocumentStore(use_bm25=False)
|
|
retriever = TfidfRetriever(document_store=ds)
|
|
reader = FARMReader(
|
|
model_name_or_path="deepset/bert-medium-squad2-distilled", use_gpu=False, top_k_per_sample=5, num_processes=0
|
|
)
|
|
ds.write_documents(docs)
|
|
|
|
base_pipeline = ExtractiveQAPipeline(reader=reader, retriever=retriever)
|
|
pipeline = TranslationWrapperPipeline(
|
|
input_translator=de_to_en_translator, output_translator=en_to_de_translator, pipeline=base_pipeline
|
|
)
|
|
|
|
prediction = pipeline.run(query="Wer lebt in Berlin?", params={"Reader": {"top_k": 3}})
|
|
assert prediction is not None
|
|
assert prediction["query"] == "Wer lebt in Berlin?"
|
|
assert "Carla" in prediction["answers"][0].answer
|
|
assert prediction["answers"][0].score <= 1
|
|
assert prediction["answers"][0].score >= 0
|
|
assert prediction["answers"][0].meta["meta_field"] == "test1"
|
|
assert prediction["answers"][0].context == "My name is Carla and I live in Berlin"
|