test_processor_layoutxlm.py 23.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# Copyright 2021 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 json
import os
import shutil
import tempfile
import unittest
from typing import List

22
23
import numpy as np

24
25
26
27
28
29
30
31
32
from transformers import PreTrainedTokenizer, PreTrainedTokenizerBase, PreTrainedTokenizerFast
from transformers.models.layoutxlm import LayoutXLMTokenizer, LayoutXLMTokenizerFast
from transformers.testing_utils import (
    require_pytesseract,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)
33
from transformers.utils import FEATURE_EXTRACTOR_NAME, cached_property, is_pytesseract_available
34
35
36
37
38


if is_pytesseract_available():
    from PIL import Image

39
    from transformers import LayoutLMv2ImageProcessor, LayoutXLMProcessor
40
41
42
43
44
45
46
47
48
49


@require_pytesseract
@require_sentencepiece
@require_tokenizers
class LayoutXLMProcessorTest(unittest.TestCase):
    tokenizer_class = LayoutXLMTokenizer
    rust_tokenizer_class = LayoutXLMTokenizerFast

    def setUp(self):
50
        image_processor_map = {
51
52
53
54
55
56
57
58
            "do_resize": True,
            "size": 224,
            "apply_ocr": True,
        }

        self.tmpdirname = tempfile.mkdtemp()
        self.feature_extraction_file = os.path.join(self.tmpdirname, FEATURE_EXTRACTOR_NAME)
        with open(self.feature_extraction_file, "w", encoding="utf-8") as fp:
59
            fp.write(json.dumps(image_processor_map) + "\n")
60

Yih-Dar's avatar
Yih-Dar committed
61
62
63
        # taken from `test_tokenization_layoutxlm.LayoutXLMTokenizationTest.test_save_pretrained`
        self.tokenizer_pretrained_name = "hf-internal-testing/tiny-random-layoutxlm"

64
    def get_tokenizer(self, **kwargs) -> PreTrainedTokenizer:
Yih-Dar's avatar
Yih-Dar committed
65
        return self.tokenizer_class.from_pretrained(self.tokenizer_pretrained_name, **kwargs)
66
67

    def get_rust_tokenizer(self, **kwargs) -> PreTrainedTokenizerFast:
Yih-Dar's avatar
Yih-Dar committed
68
        return self.rust_tokenizer_class.from_pretrained(self.tokenizer_pretrained_name, **kwargs)
69
70
71
72

    def get_tokenizers(self, **kwargs) -> List[PreTrainedTokenizerBase]:
        return [self.get_tokenizer(**kwargs), self.get_rust_tokenizer(**kwargs)]

73
74
    def get_image_processor(self, **kwargs):
        return LayoutLMv2ImageProcessor.from_pretrained(self.tmpdirname, **kwargs)
75
76
77
78

    def tearDown(self):
        shutil.rmtree(self.tmpdirname)

79
80
81
82
83
84
85
86
87
88
89
    def prepare_image_inputs(self):
        """This function prepares a list of PIL images, or a list of numpy arrays if one specifies numpify=True,
        or a list of PyTorch tensors if one specifies torchify=True.
        """

        image_inputs = [np.random.randint(255, size=(3, 30, 400), dtype=np.uint8)]

        image_inputs = [Image.fromarray(np.moveaxis(x, 0, -1)) for x in image_inputs]

        return image_inputs

90
    def test_save_load_pretrained_default(self):
91
        image_processor = self.get_image_processor()
92
93
        tokenizers = self.get_tokenizers()
        for tokenizer in tokenizers:
94
            processor = LayoutXLMProcessor(image_processor=image_processor, tokenizer=tokenizer)
95
96
97
98
99
100
101

            processor.save_pretrained(self.tmpdirname)
            processor = LayoutXLMProcessor.from_pretrained(self.tmpdirname)

            self.assertEqual(processor.tokenizer.get_vocab(), tokenizer.get_vocab())
            self.assertIsInstance(processor.tokenizer, (LayoutXLMTokenizer, LayoutXLMTokenizerFast))

102
103
            self.assertEqual(processor.image_processor.to_json_string(), image_processor.to_json_string())
            self.assertIsInstance(processor.image_processor, LayoutLMv2ImageProcessor)
104
105

    def test_save_load_pretrained_additional_features(self):
106
        processor = LayoutXLMProcessor(image_processor=self.get_image_processor(), tokenizer=self.get_tokenizer())
107
108
109
110
        processor.save_pretrained(self.tmpdirname)

        # slow tokenizer
        tokenizer_add_kwargs = self.get_tokenizer(bos_token="(BOS)", eos_token="(EOS)")
111
        image_processor_add_kwargs = self.get_image_processor(do_resize=False, size=30)
112
113
114
115
116
117
118
119
120
121
122
123
124

        processor = LayoutXLMProcessor.from_pretrained(
            self.tmpdirname,
            use_fast=False,
            bos_token="(BOS)",
            eos_token="(EOS)",
            do_resize=False,
            size=30,
        )

        self.assertEqual(processor.tokenizer.get_vocab(), tokenizer_add_kwargs.get_vocab())
        self.assertIsInstance(processor.tokenizer, LayoutXLMTokenizer)

125
126
        self.assertEqual(processor.image_processor.to_json_string(), image_processor_add_kwargs.to_json_string())
        self.assertIsInstance(processor.image_processor, LayoutLMv2ImageProcessor)
127
128
129

        # fast tokenizer
        tokenizer_add_kwargs = self.get_rust_tokenizer(bos_token="(BOS)", eos_token="(EOS)")
130
        image_processor_add_kwargs = self.get_image_processor(do_resize=False, size=30)
131
132
133
134
135
136
137
138

        processor = LayoutXLMProcessor.from_pretrained(
            self.tmpdirname, use_xlm=True, bos_token="(BOS)", eos_token="(EOS)", do_resize=False, size=30
        )

        self.assertEqual(processor.tokenizer.get_vocab(), tokenizer_add_kwargs.get_vocab())
        self.assertIsInstance(processor.tokenizer, LayoutXLMTokenizerFast)

139
140
        self.assertEqual(processor.image_processor.to_json_string(), image_processor_add_kwargs.to_json_string())
        self.assertIsInstance(processor.image_processor, LayoutLMv2ImageProcessor)
141

142
    def test_model_input_names(self):
143
        image_processor = self.get_image_processor()
144
145
        tokenizer = self.get_tokenizer()

146
        processor = LayoutXLMProcessor(tokenizer=tokenizer, image_processor=image_processor)
147
148
149
150
151
152
153
154
155

        input_str = "lower newer"
        image_input = self.prepare_image_inputs()

        # add extra args
        inputs = processor(text=input_str, images=image_input, return_codebook_pixels=False, return_image_mask=False)

        self.assertListEqual(list(inputs.keys()), processor.model_input_names)

156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
    @slow
    def test_overflowing_tokens(self):
        # In the case of overflowing tokens, test that we still have 1-to-1 mapping between the images and input_ids (sequences that are too long are broken down into multiple sequences).

        from datasets import load_dataset

        # set up
        datasets = load_dataset("nielsr/funsd")
        processor = LayoutXLMProcessor.from_pretrained("microsoft/layoutxlm-base", apply_ocr=False)

        def preprocess_data(examples):
            images = [Image.open(path).convert("RGB") for path in examples["image_path"]]
            words = examples["words"]
            boxes = examples["bboxes"]
            word_labels = examples["ner_tags"]
            encoded_inputs = processor(
                images,
                words,
                boxes=boxes,
                word_labels=word_labels,
                max_length=512,
                padding="max_length",
                truncation=True,
                return_overflowing_tokens=True,
                stride=50,
                return_offsets_mapping=True,
                return_tensors="pt",
            )
            return encoded_inputs

        train_data = preprocess_data(datasets["train"])

        self.assertEqual(len(train_data["image"]), len(train_data["input_ids"]))

190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217

# different use cases tests
@require_sentencepiece
@require_torch
@require_pytesseract
class LayoutXLMProcessorIntegrationTests(unittest.TestCase):
    @cached_property
    def get_images(self):
        # we verify our implementation on 2 document images from the DocVQA dataset
        from datasets import load_dataset

        ds = load_dataset("hf-internal-testing/fixtures_docvqa", split="test")

        image_1 = Image.open(ds[0]["file"]).convert("RGB")
        image_2 = Image.open(ds[1]["file"]).convert("RGB")

        return image_1, image_2

    @cached_property
    def get_tokenizers(self):
        slow_tokenizer = LayoutXLMTokenizer.from_pretrained("microsoft/layoutxlm-base")
        fast_tokenizer = LayoutXLMTokenizerFast.from_pretrained("microsoft/layoutxlm-base")
        return [slow_tokenizer, fast_tokenizer]

    @slow
    def test_processor_case_1(self):
        # case 1: document image classification (training, inference) + token classification (inference), apply_ocr = True

218
        image_processor = LayoutLMv2ImageProcessor()
219
220
221
222
        tokenizers = self.get_tokenizers
        images = self.get_images

        for tokenizer in tokenizers:
223
            processor = LayoutXLMProcessor(image_processor=image_processor, tokenizer=tokenizer)
224
225

            # not batched
226
            input_feat_extract = image_processor(images[0], return_tensors="pt")
227
228
229
230
            input_processor = processor(images[0], return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
231
            actual_keys = sorted(input_processor.keys())
232
233
234
235
236
237
238
239
            self.assertListEqual(actual_keys, expected_keys)

            # verify image
            self.assertAlmostEqual(
                input_feat_extract["pixel_values"].sum(), input_processor["image"].sum(), delta=1e-2
            )

            # verify input_ids
NielsRogge's avatar
NielsRogge committed
240
            # this was obtained with Tesseract 4.1.1
241
242
243
            # fmt: off
            expected_decoding = "<s> 11:14 to 11:39 a.m 11:39 to 11:44 a.m. 11:44 a.m. to 12:25 p.m. 12:25 to 12:58 p.m. 12:58 to 4:00 p.m. 2:00 to 5:00 p.m. Coffee Break Coffee will be served for men and women in the lobby adjacent to exhibit area. Please move into exhibit area. (Exhibits Open) TRRF GENERAL SESSION (PART |) Presiding: Lee A. Waller TRRF Vice President “Introductory Remarks” Lee A. Waller, TRRF Vice Presi- dent Individual Interviews with TRRF Public Board Members and Sci- entific Advisory Council Mem- bers Conducted by TRRF Treasurer Philip G. Kuehn to get answers which the public refrigerated warehousing industry is looking for. Plus questions from the floor. Dr. Emil M. Mrak, University of Cal- ifornia, Chairman, TRRF Board; Sam R. Cecil, University of Georgia College of Agriculture; Dr. Stanley Charm, Tufts University School of Medicine; Dr. Robert H. Cotton, ITT Continental Baking Company; Dr. Owen Fennema, University of Wis- consin; Dr. Robert E. Hardenburg, USDA. Questions and Answers Exhibits Open Capt. Jack Stoney Room TRRF Scientific Advisory Council Meeting Ballroom Foyer</s>"  # noqa: E231
            # fmt: on
NielsRogge's avatar
NielsRogge committed
244
            decoding = processor.decode(input_processor.input_ids.squeeze().tolist())
245
246
247
            self.assertSequenceEqual(decoding, expected_decoding)

            # batched
248
            input_feat_extract = image_processor(images, return_tensors="pt")
249
250
251
252
            input_processor = processor(images, padding=True, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
253
            actual_keys = sorted(input_processor.keys())
254
255
256
257
258
259
260
261
            self.assertListEqual(actual_keys, expected_keys)

            # verify images
            self.assertAlmostEqual(
                input_feat_extract["pixel_values"].sum(), input_processor["image"].sum(), delta=1e-2
            )

            # verify input_ids
NielsRogge's avatar
NielsRogge committed
262
            # this was obtained with Tesseract 4.1.1
263
264
265
            # fmt: off
            expected_decoding = "<s> 7 ITC Limited REPORT AND ACCOUNTS 2013 ITC’s Brands: An Asset for the Nation The consumer needs and aspirations they fulfil, the benefit they generate for millions across ITC’s value chains, the future-ready capabilities that support them, and the value that they create for the country, have made ITC’s brands national assets, adding to India’s competitiveness. It is ITC’s aspiration to be the No 1 FMCG player in the country, driven by its new FMCG businesses. A recent Nielsen report has highlighted that ITC's new FMCG businesses are the fastest growing among the top consumer goods companies operating in India. ITC takes justifiable pride that, along with generating economic value, these celebrated Indian brands also drive the creation of larger societal capital through the virtuous cycle of sustainable and inclusive growth. DI WILLS * ; LOVE DELIGHTFULLY SOFT SKIN? aia Ans Source: https://www.industrydocuments.ucsf.edu/docs/snbx0223</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>"  # noqa: E231
            # fmt: on
NielsRogge's avatar
NielsRogge committed
266
            decoding = processor.decode(input_processor.input_ids[1].tolist())
267
268
269
270
271
272
            self.assertSequenceEqual(decoding, expected_decoding)

    @slow
    def test_processor_case_2(self):
        # case 2: document image classification (training, inference) + token classification (inference), apply_ocr=False

273
        image_processor = LayoutLMv2ImageProcessor(apply_ocr=False)
274
275
276
277
        tokenizers = self.get_tokenizers
        images = self.get_images

        for tokenizer in tokenizers:
278
            processor = LayoutXLMProcessor(image_processor=image_processor, tokenizer=tokenizer)
279
280
281
282
283
284
285
286
287
288
289
290
291
292

            # not batched
            words = ["hello", "world"]
            boxes = [[1, 2, 3, 4], [5, 6, 7, 8]]
            input_processor = processor(images[0], words, boxes=boxes, return_tensors="pt")

            # verify keys
            expected_keys = ["input_ids", "bbox", "attention_mask", "image"]
            actual_keys = list(input_processor.keys())
            for key in expected_keys:
                self.assertIn(key, actual_keys)

            # verify input_ids
            expected_decoding = "<s> hello world</s>"
NielsRogge's avatar
NielsRogge committed
293
            decoding = processor.decode(input_processor.input_ids.squeeze().tolist())
294
295
296
297
298
299
300
301
302
            self.assertSequenceEqual(decoding, expected_decoding)

            # batched
            words = [["hello", "world"], ["my", "name", "is", "niels"]]
            boxes = [[[1, 2, 3, 4], [5, 6, 7, 8]], [[3, 2, 5, 1], [6, 7, 4, 2], [3, 9, 2, 4], [1, 1, 2, 3]]]
            input_processor = processor(images, words, boxes=boxes, padding=True, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
303
            actual_keys = sorted(input_processor.keys())
304
305
306
307
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "<s> hello world</s><pad><pad>"
NielsRogge's avatar
NielsRogge committed
308
            decoding = processor.decode(input_processor.input_ids[0].tolist())
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify bbox
            expected_bbox = [
                [0, 0, 0, 0],
                [3, 2, 5, 1],
                [6, 7, 4, 2],
                [3, 9, 2, 4],
                [1, 1, 2, 3],
                [1, 1, 2, 3],
                [1000, 1000, 1000, 1000],
            ]
            self.assertListEqual(input_processor.bbox[1].tolist(), expected_bbox)

    @slow
    def test_processor_case_3(self):
        # case 3: token classification (training), apply_ocr=False

327
        image_processor = LayoutLMv2ImageProcessor(apply_ocr=False)
328
329
330
331
        tokenizers = self.get_tokenizers
        images = self.get_images

        for tokenizer in tokenizers:
332
            processor = LayoutXLMProcessor(image_processor=image_processor, tokenizer=tokenizer)
333
334
335
336
337
338
339
340
341

            # not batched
            words = ["weirdly", "world"]
            boxes = [[1, 2, 3, 4], [5, 6, 7, 8]]
            word_labels = [1, 2]
            input_processor = processor(images[0], words, boxes=boxes, word_labels=word_labels, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids", "labels"]
342
            actual_keys = sorted(input_processor.keys())
343
344
345
346
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "<s> weirdly world</s>"
NielsRogge's avatar
NielsRogge committed
347
            decoding = processor.decode(input_processor.input_ids.squeeze().tolist())
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify labels
            expected_labels = [-100, 1, -100, 2, -100]
            self.assertListEqual(input_processor.labels.squeeze().tolist(), expected_labels)

            # batched
            words = [["hello", "world"], ["my", "name", "is", "niels"]]
            boxes = [[[1, 2, 3, 4], [5, 6, 7, 8]], [[3, 2, 5, 1], [6, 7, 4, 2], [3, 9, 2, 4], [1, 1, 2, 3]]]
            word_labels = [[1, 2], [6, 3, 10, 2]]
            input_processor = processor(
                images, words, boxes=boxes, word_labels=word_labels, padding=True, return_tensors="pt"
            )

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids", "labels"]
364
            actual_keys = sorted(input_processor.keys())
365
366
367
368
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "<s> my name is niels</s>"
NielsRogge's avatar
NielsRogge committed
369
            decoding = processor.decode(input_processor.input_ids[1].tolist())
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify bbox
            expected_bbox = [
                [0, 0, 0, 0],
                [3, 2, 5, 1],
                [6, 7, 4, 2],
                [3, 9, 2, 4],
                [1, 1, 2, 3],
                [1, 1, 2, 3],
                [1000, 1000, 1000, 1000],
            ]
            self.assertListEqual(input_processor.bbox[1].tolist(), expected_bbox)

            # verify labels
            expected_labels = [-100, 6, 3, 10, 2, -100, -100]
            self.assertListEqual(input_processor.labels[1].tolist(), expected_labels)

    @slow
    def test_processor_case_4(self):
        # case 4: visual question answering (inference), apply_ocr=True

392
        image_processor = LayoutLMv2ImageProcessor()
393
394
395
396
        tokenizers = self.get_tokenizers
        images = self.get_images

        for tokenizer in tokenizers:
397
            processor = LayoutXLMProcessor(image_processor=image_processor, tokenizer=tokenizer)
398
399
400
401
402
403
404

            # not batched
            question = "What's his name?"
            input_processor = processor(images[0], question, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
405
            actual_keys = sorted(input_processor.keys())
406
407
408
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
NielsRogge's avatar
NielsRogge committed
409
            # this was obtained with Tesseract 4.1.1
410
411
412
            # fmt: off
            expected_decoding = "<s> What's his name?</s></s> 11:14 to 11:39 a.m 11:39 to 11:44 a.m. 11:44 a.m. to 12:25 p.m. 12:25 to 12:58 p.m. 12:58 to 4:00 p.m. 2:00 to 5:00 p.m. Coffee Break Coffee will be served for men and women in the lobby adjacent to exhibit area. Please move into exhibit area. (Exhibits Open) TRRF GENERAL SESSION (PART |) Presiding: Lee A. Waller TRRF Vice President “Introductory Remarks” Lee A. Waller, TRRF Vice Presi- dent Individual Interviews with TRRF Public Board Members and Sci- entific Advisory Council Mem- bers Conducted by TRRF Treasurer Philip G. Kuehn to get answers which the public refrigerated warehousing industry is looking for. Plus questions from the floor. Dr. Emil M. Mrak, University of Cal- ifornia, Chairman, TRRF Board; Sam R. Cecil, University of Georgia College of Agriculture; Dr. Stanley Charm, Tufts University School of Medicine; Dr. Robert H. Cotton, ITT Continental Baking Company; Dr. Owen Fennema, University of Wis- consin; Dr. Robert E. Hardenburg, USDA. Questions and Answers Exhibits Open Capt. Jack Stoney Room TRRF Scientific Advisory Council Meeting Ballroom Foyer</s>"  # noqa: E231
            # fmt: on
NielsRogge's avatar
NielsRogge committed
413
            decoding = processor.decode(input_processor.input_ids.squeeze().tolist())
414
415
416
417
418
419
420
421
422
423
            self.assertSequenceEqual(decoding, expected_decoding)

            # batched
            questions = ["How old is he?", "what's the time"]
            input_processor = processor(
                images, questions, padding="max_length", max_length=20, truncation=True, return_tensors="pt"
            )

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
424
            actual_keys = sorted(input_processor.keys())
425
426
427
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
NielsRogge's avatar
NielsRogge committed
428
            # this was obtained with Tesseract 4.1.1
429
            expected_decoding = "<s> what's the time</s></s> 7 ITC Limited REPORT AND ACCOUNTS 2013</s>"
NielsRogge's avatar
NielsRogge committed
430
            decoding = processor.decode(input_processor.input_ids[1].tolist())
431
432
433
434
435
436
437
438
439
440
441
442
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify bbox
            # fmt: off
            expected_bbox = [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [1000, 1000, 1000, 1000], [1000, 1000, 1000, 1000], [0, 45, 67, 80], [72, 56, 109, 67], [72, 56, 109, 67], [116, 56, 189, 67], [198, 59, 253, 66], [257, 59, 285, 66], [289, 59, 365, 66], [289, 59, 365, 66], [289, 59, 365, 66], [289, 59, 365, 66], [372, 59, 407, 66], [1000, 1000, 1000, 1000]]  # noqa: E231
            # fmt: on
            self.assertListEqual(input_processor.bbox[1].tolist(), expected_bbox)

    @slow
    def test_processor_case_5(self):
        # case 5: visual question answering (inference), apply_ocr=False

443
        image_processor = LayoutLMv2ImageProcessor(apply_ocr=False)
444
445
446
447
        tokenizers = self.get_tokenizers
        images = self.get_images

        for tokenizer in tokenizers:
448
            processor = LayoutXLMProcessor(image_processor=image_processor, tokenizer=tokenizer)
449
450
451
452
453
454
455
456
457

            # not batched
            question = "What's his name?"
            words = ["hello", "world"]
            boxes = [[1, 2, 3, 4], [5, 6, 7, 8]]
            input_processor = processor(images[0], question, words, boxes, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
458
            actual_keys = sorted(input_processor.keys())
459
460
461
462
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "<s> What's his name?</s></s> hello world</s>"
NielsRogge's avatar
NielsRogge committed
463
            decoding = processor.decode(input_processor.input_ids.squeeze().tolist())
464
465
466
467
468
469
470
471
472
473
            self.assertSequenceEqual(decoding, expected_decoding)

            # batched
            questions = ["How old is he?", "what's the time"]
            words = [["hello", "world"], ["my", "name", "is", "niels"]]
            boxes = [[[1, 2, 3, 4], [5, 6, 7, 8]], [[3, 2, 5, 1], [6, 7, 4, 2], [3, 9, 2, 4], [1, 1, 2, 3]]]
            input_processor = processor(images, questions, words, boxes, padding=True, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
474
            actual_keys = sorted(input_processor.keys())
475
476
477
478
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "<s> How old is he?</s></s> hello world</s><pad><pad>"
NielsRogge's avatar
NielsRogge committed
479
            decoding = processor.decode(input_processor.input_ids[0].tolist())
480
481
482
            self.assertSequenceEqual(decoding, expected_decoding)

            expected_decoding = "<s> what's the time</s></s> my name is niels</s>"
NielsRogge's avatar
NielsRogge committed
483
            decoding = processor.decode(input_processor.input_ids[1].tolist())
484
485
486
487
488
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify bbox
            expected_bbox = [[6, 7, 4, 2], [3, 9, 2, 4], [1, 1, 2, 3], [1, 1, 2, 3], [1000, 1000, 1000, 1000]]
            self.assertListEqual(input_processor.bbox[1].tolist()[-5:], expected_bbox)