test_pipelines_document_question_answering.py 13.2 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest

from transformers import MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING, AutoTokenizer, is_vision_available
from transformers.pipelines import pipeline
from transformers.pipelines.document_question_answering import apply_tesseract
from transformers.testing_utils import (
21
    is_pipeline_test,
22
23
24
25
26
27
28
29
30
    nested_simplify,
    require_detectron2,
    require_pytesseract,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

31
from .test_pipelines_common import ANY
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55


if is_vision_available():
    from PIL import Image

    from transformers.image_utils import load_image
else:

    class Image:
        @staticmethod
        def open(*args, **kwargs):
            pass

    def load_image(_):
        return None


# This is a pinned image from a specific revision of a document question answering space, hosted by HuggingFace,
# so we can expect it to be available.
INVOICE_URL = (
    "https://huggingface.co/spaces/impira/docquery/resolve/2f6c96314dc84dfda62d40de9da55f2f5165d403/invoice.png"
)


56
@is_pipeline_test
57
58
@require_torch
@require_vision
59
class DocumentQuestionAnsweringPipelineTests(unittest.TestCase):
60
61
62
63
    model_mapping = MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING

    @require_pytesseract
    @require_vision
64
    def get_test_pipeline(self, model, tokenizer, processor, torch_dtype="float32"):
65
        dqa_pipeline = pipeline(
66
67
68
69
70
            "document-question-answering",
            model=model,
            tokenizer=tokenizer,
            image_processor=processor,
            torch_dtype=torch_dtype,
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
        )

        image = INVOICE_URL
        word_boxes = list(zip(*apply_tesseract(load_image(image), None, "")))
        question = "What is the placebo?"
        examples = [
            {
                "image": load_image(image),
                "question": question,
            },
            {
                "image": image,
                "question": question,
            },
            {
                "image": image,
                "question": question,
                "word_boxes": word_boxes,
            },
        ]
        return dqa_pipeline, examples

    def run_pipeline_test(self, dqa_pipeline, examples):
        outputs = dqa_pipeline(examples, top_k=2)
        self.assertEqual(
            outputs,
            [
                [
                    {"score": ANY(float), "answer": ANY(str), "start": ANY(int), "end": ANY(int)},
                    {"score": ANY(float), "answer": ANY(str), "start": ANY(int), "end": ANY(int)},
                ]
            ]
Yih-Dar's avatar
Yih-Dar committed
103
            * 3,
104
105
106
107
108
109
        )

    @require_torch
    @require_detectron2
    @require_pytesseract
    def test_small_model_pt(self):
110
111
112
        dqa_pipeline = pipeline(
            "document-question-answering", model="hf-internal-testing/tiny-random-layoutlmv2-for-dqa-test"
        )
113
114
115
116
        image = INVOICE_URL
        question = "How many cats are there?"

        expected_output = [
117
118
            {"score": 0.0001, "answer": "oy 2312/2019", "start": 38, "end": 39},
            {"score": 0.0001, "answer": "oy 2312/2019 DUE", "start": 38, "end": 40},
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
        ]
        outputs = dqa_pipeline(image=image, question=question, top_k=2)
        self.assertEqual(nested_simplify(outputs, decimals=4), expected_output)

        outputs = dqa_pipeline({"image": image, "question": question}, top_k=2)
        self.assertEqual(nested_simplify(outputs, decimals=4), expected_output)

        # This image does not detect ANY text in it, meaning layoutlmv2 should fail.
        # Empty answer probably
        image = "./tests/fixtures/tests_samples/COCO/000000039769.png"
        outputs = dqa_pipeline(image=image, question=question, top_k=2)
        self.assertEqual(outputs, [])

        # We can optionnally pass directly the words and bounding boxes
        image = "./tests/fixtures/tests_samples/COCO/000000039769.png"
        words = []
        boxes = []
        outputs = dqa_pipeline(image=image, question=question, words=words, boxes=boxes, top_k=2)
        self.assertEqual(outputs, [])

    # 	 TODO: Enable this once hf-internal-testing/tiny-random-donut is implemented
    #    @require_torch
    #    def test_small_model_pt_donut(self):
    #        dqa_pipeline = pipeline("document-question-answering", model="hf-internal-testing/tiny-random-donut")
    #        # dqa_pipeline = pipeline("document-question-answering", model="../tiny-random-donut")
    #        image = "https://templates.invoicehome.com/invoice-template-us-neat-750px.png"
    #        question = "How many cats are there?"
    #
    #        outputs = dqa_pipeline(image=image, question=question, top_k=2)
    #        self.assertEqual(
    #            nested_simplify(outputs, decimals=4), [{"score": 0.8799, "answer": "2"}, {"score": 0.296, "answer": "1"}]
    #        )

    @slow
    @require_torch
    @require_detectron2
    @require_pytesseract
    def test_large_model_pt(self):
        dqa_pipeline = pipeline(
            "document-question-answering",
            model="tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa",
            revision="9977165",
        )
        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),
            [
169
170
                {"score": 0.9944, "answer": "us-001", "start": 16, "end": 16},
                {"score": 0.0009, "answer": "us-001", "start": 16, "end": 16},
171
172
173
174
175
176
177
            ],
        )

        outputs = dqa_pipeline({"image": image, "question": question}, top_k=2)
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
178
179
                {"score": 0.9944, "answer": "us-001", "start": 16, "end": 16},
                {"score": 0.0009, "answer": "us-001", "start": 16, "end": 16},
180
181
182
183
184
185
186
187
188
189
            ],
        )

        outputs = dqa_pipeline(
            [{"image": image, "question": question}, {"image": image, "question": question}], top_k=2
        )
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
                [
190
191
                    {"score": 0.9944, "answer": "us-001", "start": 16, "end": 16},
                    {"score": 0.0009, "answer": "us-001", "start": 16, "end": 16},
192
193
194
195
196
                ],
            ]
            * 2,
        )

197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
    @slow
    @require_torch
    @require_detectron2
    @require_pytesseract
    def test_large_model_pt_chunk(self):
        dqa_pipeline = pipeline(
            "document-question-answering",
            model="tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa",
            revision="9977165",
            max_seq_len=50,
        )
        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),
            [
215
216
                {"score": 0.9974, "answer": "1110212019", "start": 23, "end": 23},
                {"score": 0.9948, "answer": "us-001", "start": 16, "end": 16},
217
218
219
220
221
222
223
            ],
        )

        outputs = dqa_pipeline({"image": image, "question": question}, top_k=2)
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
224
225
                {"score": 0.9974, "answer": "1110212019", "start": 23, "end": 23},
                {"score": 0.9948, "answer": "us-001", "start": 16, "end": 16},
226
227
228
229
230
231
232
233
234
235
            ],
        )

        outputs = dqa_pipeline(
            [{"image": image, "question": question}, {"image": image, "question": question}], top_k=2
        )
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
                [
236
237
                    {"score": 0.9974, "answer": "1110212019", "start": 23, "end": 23},
                    {"score": 0.9948, "answer": "us-001", "start": 16, "end": 16},
238
239
240
241
242
                ]
            ]
            * 2,
        )

243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
    @slow
    @require_torch
    @require_pytesseract
    @require_vision
    def test_large_model_pt_layoutlm(self):
        tokenizer = AutoTokenizer.from_pretrained(
            "impira/layoutlm-document-qa", revision="3dc6de3", add_prefix_space=True
        )
        dqa_pipeline = pipeline(
            "document-question-answering",
            model="impira/layoutlm-document-qa",
            tokenizer=tokenizer,
            revision="3dc6de3",
        )
        image = INVOICE_URL
        question = "What is the invoice number?"

        outputs = dqa_pipeline(image=image, question=question, top_k=2)
        self.assertEqual(
262
            nested_simplify(outputs, decimals=3),
263
            [
264
265
                {"score": 0.425, "answer": "us-001", "start": 16, "end": 16},
                {"score": 0.082, "answer": "1110212019", "start": 23, "end": 23},
266
267
268
269
270
            ],
        )

        outputs = dqa_pipeline({"image": image, "question": question}, top_k=2)
        self.assertEqual(
271
            nested_simplify(outputs, decimals=3),
272
            [
273
274
                {"score": 0.425, "answer": "us-001", "start": 16, "end": 16},
                {"score": 0.082, "answer": "1110212019", "start": 23, "end": 23},
275
276
277
278
279
280
281
            ],
        )

        outputs = dqa_pipeline(
            [{"image": image, "question": question}, {"image": image, "question": question}], top_k=2
        )
        self.assertEqual(
282
            nested_simplify(outputs, decimals=3),
283
284
            [
                [
285
286
                    {"score": 0.425, "answer": "us-001", "start": 16, "end": 16},
                    {"score": 0.082, "answer": "1110212019", "start": 23, "end": 23},
287
288
289
290
291
292
293
294
295
296
                ]
            ]
            * 2,
        )

        word_boxes = list(zip(*apply_tesseract(load_image(image), None, "")))

        # This model should also work if `image` is set to None
        outputs = dqa_pipeline({"image": None, "word_boxes": word_boxes, "question": question}, top_k=2)
        self.assertEqual(
297
            nested_simplify(outputs, decimals=3),
298
            [
299
300
                {"score": 0.425, "answer": "us-001", "start": 16, "end": 16},
                {"score": 0.082, "answer": "1110212019", "start": 23, "end": 23},
301
302
303
            ],
        )

304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
    @slow
    @require_torch
    @require_pytesseract
    @require_vision
    def test_large_model_pt_layoutlm_chunk(self):
        tokenizer = AutoTokenizer.from_pretrained(
            "impira/layoutlm-document-qa", revision="3dc6de3", add_prefix_space=True
        )
        dqa_pipeline = pipeline(
            "document-question-answering",
            model="impira/layoutlm-document-qa",
            tokenizer=tokenizer,
            revision="3dc6de3",
            max_seq_len=50,
        )
        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),
            [
326
327
                {"score": 0.9999, "answer": "us-001", "start": 16, "end": 16},
                {"score": 0.9998, "answer": "us-001", "start": 16, "end": 16},
328
329
330
331
332
333
334
335
336
337
            ],
        )

        outputs = dqa_pipeline(
            [{"image": image, "question": question}, {"image": image, "question": question}], top_k=2
        )
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
                [
338
339
                    {"score": 0.9999, "answer": "us-001", "start": 16, "end": 16},
                    {"score": 0.9998, "answer": "us-001", "start": 16, "end": 16},
340
341
342
343
344
345
346
347
348
349
350
351
                ]
            ]
            * 2,
        )

        word_boxes = list(zip(*apply_tesseract(load_image(image), None, "")))

        # This model should also work if `image` is set to None
        outputs = dqa_pipeline({"image": None, "word_boxes": word_boxes, "question": question}, top_k=2)
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
352
353
                {"score": 0.9999, "answer": "us-001", "start": 16, "end": 16},
                {"score": 0.9998, "answer": "us-001", "start": 16, "end": 16},
354
355
356
            ],
        )

357
358
359
360
361
362
363
    @slow
    @require_torch
    def test_large_model_pt_donut(self):
        dqa_pipeline = pipeline(
            "document-question-answering",
            model="naver-clova-ix/donut-base-finetuned-docvqa",
            tokenizer=AutoTokenizer.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa"),
364
            image_processor="naver-clova-ix/donut-base-finetuned-docvqa",
365
366
367
368
369
        )

        image = INVOICE_URL
        question = "What is the invoice number?"
        outputs = dqa_pipeline(image=image, question=question, top_k=2)
370
        self.assertEqual(nested_simplify(outputs, decimals=4), [{"answer": "us-001"}])
371
372

    @require_tf
amyeroberts's avatar
amyeroberts committed
373
    @unittest.skip(reason="Document question answering not implemented in TF")
374
375
    def test_small_model_tf(self):
        pass