dnrops
/
transformers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Change document question answering pipeline to always return an array (#19071) 

Co-authored-by: Ankur Goyal <ankur@impira.com>
This commit is contained in:
Ankur Goyal 2022-09-20 06:17:57 -07:00 committed by GitHub
parent cc567e0063
commit 67403413bd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 1 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/pipelines/document_question_answering.py
View File

@ -383,8 +383,6 @@ class DocumentQuestionAnsweringPipeline(Pipeline):
answers = self.postprocess_extractive_qa(model_outputs, top_k=top_k, **kwargs)
answers = sorted(answers, key=lambda x: x.get("score", 0), reverse=True)[:top_k]
if len(answers) == 1:
return answers[0]
return answers
def postprocess_donut(self, model_outputs, **kwargs):

2
tests/pipelines/test_pipelines_document_question_answering.py
View File

@ -267,7 +267,7 @@ class DocumentQuestionAnsweringPipelineTests(unittest.TestCase, metaclass=Pipeli
image = INVOICE_URL
question = "What is the invoice number?"
outputs = dqa_pipeline(image=image, question=question, top_k=2)
self.assertEqual(nested_simplify(outputs, decimals=4), {"answer": "us-001"})
self.assertEqual(nested_simplify(outputs, decimals=4), [{"answer": "us-001"}])
@require_tf
@unittest.skip("Document question answering not implemented in TF")