transformers/tests/models/kosmos2_5/test_modeling_kosmos2_5.py

# Copyright 2024 Microsoft Research and The HuggingFace Inc. 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.
"""Testing suite for the PyTorch KOSMOS-2.5 model."""

import copy
import inspect
import tempfile
import unittest

import numpy as np
import pytest
import requests
from parameterized import parameterized

from transformers import AutoProcessor, Kosmos2_5Config
from transformers.models.kosmos2_5.configuration_kosmos2_5 import (
    Kosmos2_5TextConfig,
    Kosmos2_5VisionConfig,
)
from transformers.testing_utils import (
    Expectations,
    require_flash_attn,
    require_torch,
    require_torch_accelerator,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_torch_available, is_vision_available

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import (
    ModelTesterMixin,
    floats_tensor,
    ids_tensor,
    random_attention_mask,
)
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import Kosmos2_5ForConditionalGeneration, Kosmos2_5Model


if is_vision_available():
    from PIL import Image


class Kosmos2_5VisionModelTester:
    def __init__(
        self,
        parent,
        batch_size=6,
        image_size=32,
        patch_size=4,
        num_channels=3,
        is_training=True,
        hidden_size=32,
        intermediate_size=64,
        num_hidden_layers=2,
        num_attention_heads=4,
        dropout=0.0,
        attention_dropout=0.0,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_channels = num_channels
        self.is_training = is_training
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.patch_embed_hidden_size = patch_size * patch_size * num_channels
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        self.scope = scope

        # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
        num_patches = (image_size // patch_size) ** 2
        self.seq_length = num_patches + 1

    def prepare_config_and_inputs(self):
        flattened_patches = floats_tensor([self.batch_size, self.seq_length, self.patch_embed_hidden_size + 2])
        config = self.get_config()

        return config, flattened_patches

    def get_config(self):
        return Kosmos2_5VisionConfig(
            image_size=self.image_size,
            patch_size=self.patch_size,
            num_channels=self.num_channels,
            hidden_size=self.hidden_size,
            intermediate_size=self.intermediate_size,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            patch_embed_hidden_size=self.patch_embed_hidden_size,
            dropout=self.dropout,
            attention_dropout=self.attention_dropout,
        )

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        config, flattened_patches = config_and_inputs
        inputs_dict = {"flattened_patches": flattened_patches}
        return config, inputs_dict


class Kosmos2_5TextModelTester:
    def __init__(
        self,
        parent,
        batch_size=6,
        seq_length=7,
        is_training=True,
        use_input_mask=True,
        use_labels=True,
        vocab_size=99,
        hidden_size=32,
        ffn_dim=64,
        num_hidden_layers=2,
        num_attention_heads=4,
        dropout=0.0,
        attention_dropout=0.0,
        max_position_embeddings=512,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.use_input_mask = use_input_mask
        self.use_labels = use_labels
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.ffn_dim = ffn_dim
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        self.max_position_embeddings = max_position_embeddings
        self.scope = scope

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)

        input_mask = None
        if self.use_input_mask:
            input_mask = random_attention_mask([self.batch_size, self.seq_length])

        if input_mask is not None:
            batch_size, seq_length = input_mask.shape
            rnd_start_indices = np.random.randint(1, seq_length - 1, size=(batch_size,))
            for batch_idx, start_index in enumerate(rnd_start_indices):
                input_mask[batch_idx, :start_index] = 1
                input_mask[batch_idx, start_index:] = 0

        config = self.get_config()

        return config, input_ids, input_mask

    def get_config(self):
        return Kosmos2_5TextConfig(
            vocab_size=self.vocab_size,
            embed_dim=self.hidden_size,
            ffn_dim=self.ffn_dim,
            layers=self.num_hidden_layers,
            attention_heads=self.num_attention_heads,
            dropout=self.dropout,
            attention_dropout=self.attention_dropout,
            max_position_embeddings=self.max_position_embeddings,
        )

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        config, input_ids, input_mask = config_and_inputs
        inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask}
        return config, inputs_dict


class Kosmos2_5ModelTester:
    def __init__(
        self,
        parent,
        text_kwargs=None,
        vision_kwargs=None,
        latent_query_num=3,
        is_training=True,
    ):
        if text_kwargs is None:
            text_kwargs = {}
        if vision_kwargs is None:
            vision_kwargs = {}

        self.parent = parent
        self.text_model_tester = Kosmos2_5TextModelTester(parent, **text_kwargs)
        self.vision_model_tester = Kosmos2_5VisionModelTester(parent, **vision_kwargs)
        self.batch_size = self.text_model_tester.batch_size  # need bs for batching_equivalence test
        self.seq_length = self.text_model_tester.seq_length
        self.latent_query_num = latent_query_num
        self.is_training = is_training

    def prepare_config_and_inputs(self):
        text_config, input_ids, attention_mask = self.text_model_tester.prepare_config_and_inputs()
        vision_config, flattened_patches = self.vision_model_tester.prepare_config_and_inputs()

        # build `image_embeds_position_mask`
        image_embeds_position_mask = torch.zeros_like(input_ids)
        image_embeds_position_mask[:, 1 : 1 + self.latent_query_num :] = 1

        config = self.get_config()

        return (
            config,
            input_ids,
            attention_mask,
            image_embeds_position_mask,
            flattened_patches,
        )

    def get_config(self):
        return Kosmos2_5Config(
            text_config=self.text_model_tester.get_config().to_dict(),
            vision_config=self.vision_model_tester.get_config().to_dict(),
            latent_query_num=self.latent_query_num,
        )

    def create_and_check_model(
        self,
        config,
        input_ids,
        attention_mask,
        image_embeds_position_mask,
        flattened_patches,
    ):
        model = Kosmos2_5Model(config).to(torch_device).eval()
        with torch.no_grad():
            result = model(input_ids, flattened_patches, image_embeds_position_mask, attention_mask)
        self.parent.assertEqual(
            result.last_hidden_state.shape,
            (
                self.text_model_tester.batch_size,
                self.text_model_tester.seq_length,
                self.text_model_tester.hidden_size,
            ),
        )
        self.parent.assertEqual(
            result.image_embeds.shape,
            (
                self.text_model_tester.batch_size,
                self.latent_query_num,
                self.text_model_tester.hidden_size,
            ),
        )

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        (
            config,
            input_ids,
            attention_mask,
            image_embeds_position_mask,
            flattened_patches,
        ) = config_and_inputs
        inputs_dict = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "image_embeds_position_mask": image_embeds_position_mask,
            "flattened_patches": flattened_patches,
        }
        return config, inputs_dict


@require_torch
class Kosmos2_5ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (Kosmos2_5Model, Kosmos2_5ForConditionalGeneration) if is_torch_available() else ()
    all_generative_model_classes = (Kosmos2_5ForConditionalGeneration,) if is_torch_available() else ()
    pipeline_model_mapping = (
        {
            "feature-extraction": Kosmos2_5Model,
        }
        if is_torch_available()
        else {}
    )

    test_resize_embeddings = False
    test_attention_outputs = False
    _is_composite = True

    def is_pipeline_test_to_skip(
        self,
        pipeline_test_casse_name,
        config_class,
        model_architecture,
        tokenizer_name,
        processor_name,
    ):
        return pipeline_test_casse_name == "ImageToTextPipelineTests"

    def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
        inputs_dict = copy.deepcopy(inputs_dict)

        if return_labels:
            if model_class.__name__ == "Kosmos2_5ForConditionalGeneration":
                inputs_dict["labels"] = torch.zeros(
                    (
                        self.model_tester.text_model_tester.batch_size,
                        self.model_tester.text_model_tester.seq_length,
                    ),
                    dtype=torch.long,
                    device=torch_device,
                )

        if model_class.__name__ in [
            "Kosmos2_5Model",
            "Kosmos2_5ForConditionalGeneration",
        ]:
            bs, _ = inputs_dict["input_ids"].shape
            seqlen = self.model_tester.text_model_tester.seq_length
            inputs_dict["input_ids"] = torch.arange(seqlen, device=torch_device).unsqueeze(0).expand(bs, seqlen)
            inputs_dict["input_ids"] = inputs_dict["input_ids"] % self.model_tester.text_model_tester.vocab_size
            inputs_dict["attention_mask"] = torch.ones((bs, seqlen), device=torch_device)
            inputs_dict["image_embeds_position_mask"] = torch.zeros((bs, seqlen), device=torch_device)
            inputs_dict["image_embeds_position_mask"][:, : self.model_tester.latent_query_num] = 1
        return inputs_dict

    def setUp(self):
        self.model_tester = Kosmos2_5ModelTester(self)
        self.config_tester = ConfigTester(self, config_class=Kosmos2_5Config, hidden_size=32)

    @unittest.skip("KOSMOS-2.5 doesn't support padding")
    def test_eager_padding_matches_padding_free_with_position_ids(self):
        pass

    @unittest.skip("KOSMOS-2.5 doesn't support padding")
    def test_sdpa_padding_matches_padding_free_with_position_ids(self):
        pass

    @parameterized.expand([("random",), ("same",)])
    @pytest.mark.generate
    @unittest.skip(
        "Kosmos-2.5 doesn't support assisted generation due to the need to extend `image_embeds_position_mask` length."
    )
    def test_assisted_decoding_matches_greedy_search(self):
        pass

    @pytest.mark.generate
    @unittest.skip(
        "Kosmos-2.5 doesn't support assisted generation due to the need to extend `image_embeds_position_mask` length."
    )
    def test_assisted_decoding_sample(self):
        pass

    @unittest.skip(
        "Kosmos-2.5 doesn't support assisted generation due to the need to extend `image_embeds_position_mask` length."
    )
    def test_prompt_lookup_decoding_matches_greedy_search(self):
        pass

    @unittest.skip(reason="Kosmos2-3 has no separate base model without a head.")
    def test_model_base_model_prefix(self):
        pass

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    def test_forward_signature(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            signature = inspect.signature(model.forward)
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            arg_names = [*signature.parameters.keys()]

            expected_arg_names = ["input_ids"]
            self.assertListEqual(arg_names[:1], expected_arg_names)

    def test_load_save_without_tied_weights(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        config.text_config.tie_word_embeddings = False
        for model_class in self.all_model_classes:
            model = model_class(config)
            with tempfile.TemporaryDirectory() as d:
                model.save_pretrained(d)

                model_reloaded, infos = model_class.from_pretrained(d, output_loading_info=True)
                # Checking the state dicts are correct
                reloaded_state = model_reloaded.state_dict()
                for k, v in model.state_dict().items():
                    self.assertIn(k, reloaded_state, f"Key {k} is missing from reloaded")
                    torch.testing.assert_close(
                        v,
                        reloaded_state[k],
                        msg=lambda x: f"{model_class.__name__}: Tensor {k}: {x}",
                    )
                # Checking there was no complain of missing weights
                self.assertEqual(infos["missing_keys"], set())

    # overwrite from common in order to use `self.model_tester.text_model_tester.num_hidden_layers`
    def test_hidden_states_output(self):
        def check_hidden_states_output(inputs_dict, config, model_class):
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))

            hidden_states = outputs.hidden_states

            expected_num_layers = getattr(
                self.model_tester,
                "expected_num_hidden_layers",
                self.model_tester.text_model_tester.num_hidden_layers + 1,
            )
            self.assertEqual(len(hidden_states), expected_num_layers)

            seq_length = self.model_tester.text_model_tester.seq_length

            self.assertListEqual(
                list(hidden_states[0].shape[-2:]),
                [seq_length, self.model_tester.text_model_tester.hidden_size],
            )

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            inputs_dict["output_hidden_states"] = True
            check_hidden_states_output(inputs_dict, config, model_class)

            # check that output_hidden_states also work using config
            del inputs_dict["output_hidden_states"]
            self._set_subconfig_attributes(config, "output_hidden_states", True)

            check_hidden_states_output(inputs_dict, config, model_class)

    @slow
    def test_model_from_pretrained(self):
        model_name = "microsoft/kosmos-2.5"
        model = Kosmos2_5Model.from_pretrained(model_name)
        self.assertIsNotNone(model)

    @unittest.skip(reason="Does not work on the tiny model as we keep hitting edge cases.")
    def test_model_parallelism(self):
        pass

    # TODO: ydshieh
    @require_torch_accelerator
    @slow
    @unittest.skip(reason="_update_causal_mask is not implemented yet which fails this test")
    def test_sdpa_can_dispatch_on_flash(self):
        pass

    # TODO: vasqu
    @unittest.skip(reason="why the heck does this have bigger tols")
    def test_eager_matches_sdpa_inference_24_fp32_pad_left_output_attentions(self):
        pass

    # TODO: ydshieh
    @unittest.skip(reason=" the model hasn't been added to auto class")
    def test_flash_attn_2_from_config(self):
        pass

    @unittest.skip("This test is currently not well designed for multimodal model (float type as an input).")
    def test_flash_attn_2_fp32_ln(self):
        pass

    @unittest.skip("This test is currently not well designed for multimodal model (float type as an input).")
    def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self):
        pass

    @unittest.skip("Kosmos 2.5 is multimodel and has specific input shapes.")
    def test_flash_attn_2_generate_reuse_cache(self):
        pass

    @pytest.mark.generate
    def test_generate_with_cache_matches_no_cache(self):
        """Verify that greedy generation with cache produces the same token IDs as without cache"""
        config, inputs_dict = self.prepare_config_and_inputs_for_generate()
        model = Kosmos2_5ForConditionalGeneration(config).to(torch_device).eval()

        with torch.no_grad():
            output_no_cache = model.generate(**inputs_dict, use_cache=False, max_new_tokens=5, do_sample=False)
            output_with_cache = model.generate(**inputs_dict, use_cache=True, max_new_tokens=5, do_sample=False)

        self.assertEqual(output_no_cache.tolist(), output_with_cache.tolist())

    @pytest.mark.generate
    @parameterized.expand([("greedy", 1), ("beam search", 2)])
    @unittest.skip(
        "KOSMOS-2.5 doesn't support inputs embeds. The test isn't skipped by checking input args because KOSMOS-2 has `generate()` overwritten",
    )
    def test_generate_from_inputs_embeds(self):
        pass

    @pytest.mark.generate
    def test_left_padding_compatibility(self):
        # Overwrite -- Kosmos-2.5 needs to prepare `image_embeds_position_mask`, and it must be padded accordingly
        _, inputs_dict = self.prepare_config_and_inputs_for_generate()
        input_ids = inputs_dict["input_ids"]

        def _prepare_image_embeds_position_mask(input_ids, pad_size):
            image_embeds_position_mask = torch.zeros(
                input_ids.shape[0], input_ids.shape[1] + pad_size, device=torch_device, dtype=input_ids.dtype
            )
            image_embeds_position_mask[:, (pad_size + 1) : pad_size + 1 + self.model_tester.latent_query_num] = 1
            return image_embeds_position_mask

        # `image_embeds_position_mask` is randomly generated in `prepare_config_and_inputs_for_generate`, and it must
        # match its padded version for the test to be valid -- we need to pass both
        unpadded_custom_inputs = {"image_embeds_position_mask": _prepare_image_embeds_position_mask(input_ids, 0)}
        padded_custom_inputs = {"image_embeds_position_mask": _prepare_image_embeds_position_mask(input_ids, 32)}
        super().test_left_padding_compatibility(
            unpadded_custom_inputs=unpadded_custom_inputs, padded_custom_inputs=padded_custom_inputs
        )


@require_vision
@require_torch
@slow
class Kosmos2_5ModelIntegrationTest(unittest.TestCase):
    def run_example(self, prompt, image, model, processor):
        inputs = processor(text=prompt, images=image, return_tensors="pt")
        inputs = {k: v.to(torch_device) if v is not None else None for k, v in inputs.items()}
        inputs["flattened_patches"] = inputs["flattened_patches"].to(model.dtype)

        generation_outputs = model.generate(
            **inputs,
            max_new_tokens=1024,
        )
        generated_ids = generation_outputs
        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)

        return generated_ids, generated_text

    def test_eager(self):
        url = "https://huggingface.co/microsoft/kosmos-2.5/resolve/main/receipt_00008.png"
        image = Image.open(requests.get(url, stream=True).raw)

        dtype = torch.bfloat16
        repo = "microsoft/kosmos-2.5"
        model = Kosmos2_5ForConditionalGeneration.from_pretrained(
            repo, device_map=torch_device, dtype=dtype, attn_implementation="eager"
        )
        processor = AutoProcessor.from_pretrained(repo)
        prompt = "<ocr>"
        generated_ids, generated_text = self.run_example(prompt, image, model, processor)
        EXPECTED_TEXT = Expectations(
            {
                ("cuda", 8): [
                    "<bbox><x_53><y_573><x_69><y_606></bbox>1\n<bbox><x_79><y_573><x_464><y_611></bbox>[REG] BLACK SAKURA\n<bbox><x_690><y_569><x_810><y_606></bbox>45,455\n<bbox><x_53><y_614><x_69><y_648></bbox>1\n<bbox><x_79><y_614><x_468><y_651></bbox>COOKIE DOH SAUCES\n<bbox><x_788><y_609><x_812><y_642></bbox>0\n<bbox><x_50><y_658><x_69><y_693></bbox>1\n<bbox><x_79><y_658><x_358><y_693></bbox>NATA DE COCO\n<bbox><x_790><y_652><x_814><y_683></bbox>0\n<bbox><x_31><y_742><x_820><y_781></bbox>Sub Total 45,455\n<bbox><x_27><y_781><x_822><y_827></bbox>PB1 (10%) 4,545\n<bbox><x_27><y_826><x_824><y_872></bbox>Rounding 0\n<bbox><x_24><y_872><x_827><y_921></bbox>Total 50,000\n<bbox><x_17><y_1056><x_836><y_1108></bbox>Card Payment 50,000\n"
                ],
                ("xpu", None): [
                    "<bbox><x_53><y_573><x_69><y_606></bbox>1\n<bbox><x_79><y_573><x_464><y_611></bbox>[REG] BLACK SAKURA\n<bbox><x_690><y_569><x_810><y_606></bbox>45,455\n<bbox><x_53><y_614><x_69><y_648></bbox>1\n<bbox><x_79><y_614><x_468><y_650></bbox>COOKIE DOH SAUCES\n<bbox><x_788><y_609><x_812><y_644></bbox>0\n<bbox><x_50><y_658><x_69><y_693></bbox>1\n<bbox><x_79><y_658><x_358><y_693></bbox>NATA DE COCO\n<bbox><x_790><y_652><x_814><y_687></bbox>0\n<bbox><x_31><y_742><x_820><y_781></bbox>Sub Total 45,455\n<bbox><x_27><y_781><x_822><y_827></bbox>PB1 (10%) 4,545\n<bbox><x_27><y_826><x_824><y_872></bbox>Rounding 0\n<bbox><x_24><y_872><x_827><y_921></bbox>Total 50,000\n<bbox><x_17><y_1056><x_836><y_1108></bbox>Card Payment 50,000\n"
                ],
            }
        ).get_expectation()

        self.assertListEqual(generated_text, EXPECTED_TEXT)

        prompt = "<md>"
        generated_ids, generated_text = self.run_example(prompt, image, model, processor)

        EXPECTED_TEXT = Expectations(
            {
                ("cuda", 8): [
                    "- **1 \\[REG\\] BLACK SAKURA** 45,455\n- **1 COOKIE DOH SAUCES** 0\n- **1 NATA DE COCO** 0\n- **Sub Total** 45,455\n- **PB1 (10%)** 4,545\n- **Rounding** 0\n- **Total** **50,000**\n\nCard Payment 50,000"
                ],
                ("xpu", None): [
                    "- **1 \\[REG\\] BLACK SAKURA** 45,455\n- **1 COOKIE DOH SAUCES** 0\n- **1 NATA DE COCO** 0\n- **Sub Total** 45,455\n- **PB1 (10%)** 4,545\n- **Rounding** 0\n- **Total** **50,000**\n\nCard Payment 50,000"
                ],
            }
        ).get_expectation()

        self.assertListEqual(generated_text, EXPECTED_TEXT)

    def test_sdpa(self):
        url = "https://huggingface.co/microsoft/kosmos-2.5/resolve/main/receipt_00008.png"
        image = Image.open(requests.get(url, stream=True).raw)

        dtype = torch.bfloat16
        repo = "microsoft/kosmos-2.5"
        model = Kosmos2_5ForConditionalGeneration.from_pretrained(
            repo, device_map=torch_device, dtype=dtype, attn_implementation="sdpa"
        )
        processor = AutoProcessor.from_pretrained(repo)
        prompt = "<ocr>"
        generated_ids, generated_text = self.run_example(prompt, image, model, processor)
        EXPECTED_TEXT = Expectations(
            {
                ("cuda", 7): [
                    "<bbox><x_53><y_573><x_69><y_606></bbox>1\n<bbox><x_79><y_573><x_464><y_611></bbox>[REG] BLACK SAKURA\n<bbox><x_690><y_569><x_810><y_606></bbox>45,455\n<bbox><x_53><y_614><x_69><y_648></bbox>1\n<bbox><x_79><y_614><x_468><y_651></bbox>COOKIE DOH SAUCES\n<bbox><x_788><y_609><x_812><y_642></bbox>0\n<bbox><x_50><y_658><x_69><y_693></bbox>1\n<bbox><x_79><y_658><x_358><y_693></bbox>NATA DE COCO\n<bbox><x_790><y_652><x_814><y_683></bbox>0\n<bbox><x_31><y_742><x_820><y_781></bbox>Sub Total 45,455\n<bbox><x_27><y_781><x_822><y_827></bbox>PB1 (10%) 4,545\n<bbox><x_27><y_826><x_824><y_872></bbox>Rounding 0\n<bbox><x_24><y_872><x_827><y_921></bbox>Total 50,000\n<bbox><x_17><y_1056><x_836><y_1108></bbox>Card Payment 50,000\n",
                ],
                ("cuda", 8): [
                    "<bbox><x_53><y_573><x_69><y_606></bbox>1\n<bbox><x_79><y_573><x_464><y_611></bbox>[REG] BLACK SAKURA\n<bbox><x_690><y_569><x_810><y_606></bbox>45,455\n<bbox><x_53><y_614><x_69><y_648></bbox>1\n<bbox><x_79><y_614><x_468><y_651></bbox>COOKIE DOH SAUCES\n<bbox><x_788><y_609><x_812><y_642></bbox>0\n<bbox><x_50><y_658><x_69><y_693></bbox>1\n<bbox><x_79><y_658><x_358><y_693></bbox>NATA DE COCO\n<bbox><x_790><y_652><x_814><y_683></bbox>0\n<bbox><x_31><y_742><x_820><y_781></bbox>Sub Total 45,455\n<bbox><x_27><y_781><x_822><y_827></bbox>PB1 (10%) 4,545\n<bbox><x_27><y_826><x_824><y_872></bbox>Rounding 0\n<bbox><x_24><y_872><x_827><y_921></bbox>Total 50,000\n<bbox><x_17><y_1056><x_836><y_1108></bbox>Card Payment 50,000\n"
                ],
                ("xpu", None): [
                    "<bbox><x_53><y_573><x_69><y_606></bbox>1\n<bbox><x_79><y_573><x_464><y_611></bbox>[REG] BLACK SAKURA\n<bbox><x_690><y_569><x_810><y_606></bbox>45,455\n<bbox><x_53><y_614><x_69><y_648></bbox>1\n<bbox><x_79><y_614><x_468><y_651></bbox>COOKIE DOH SAUCES\n<bbox><x_788><y_609><x_812><y_642></bbox>0\n<bbox><x_50><y_658><x_69><y_693></bbox>1\n<bbox><x_79><y_658><x_358><y_693></bbox>NATA DE COCO\n<bbox><x_790><y_652><x_814><y_683></bbox>0\n<bbox><x_31><y_742><x_820><y_781></bbox>Sub Total 45,455\n<bbox><x_27><y_781><x_822><y_827></bbox>PB1 (10%) 4,545\n<bbox><x_27><y_826><x_824><y_872></bbox>Rounding 0\n<bbox><x_24><y_872><x_827><y_921></bbox>Total 50,000\n<bbox><x_17><y_1056><x_836><y_1108></bbox>Card Payment 50,000\n"
                ],
            }
        ).get_expectation()

        self.assertListEqual(generated_text, EXPECTED_TEXT)

        prompt = "<md>"
        generated_ids, generated_text = self.run_example(prompt, image, model, processor)

        EXPECTED_TEXT = Expectations(
            {
                ("cuda", 7): [
                    "- **1 \\[REG\\] BLACK SAKURA** 45,455\n- **1 COOKIE DOH SAUCES** 0\n- **1 NATA DE COCO** 0\n- **Sub Total** 45,455\n- **PB1 (10%)** 4,545\n- **Rounding** 0\n- **Total** **50,000**\n\nCard Payment 50,000"
                ],
                ("cuda", 8): [
                    "- **1 \\[REG\\] BLACK SAKURA** 45,455\n- **1 COOKIE DOH SAUCES** 0\n- **1 NATA DE COCO** 0\n- **Sub Total** 45,455\n- **PB1 (10%)** 4,545\n- **Rounding** 0\n- **Total** **50,000**\n\nCard Payment 50,000"
                ],
                ("xpu", None): [
                    "- **1 \\[REG\\] BLACK SAKURA** 45,455\n- **1 COOKIE DOH SAUCES** 0\n- **1 NATA DE COCO** 0\n- **Sub Total** 45,455\n- **PB1 (10%)** 4,545\n- **Rounding** 0\n- **Total** **50,000**\n\nCard Payment 50,000"
                ],
            }
        ).get_expectation()

        self.assertListEqual(generated_text, EXPECTED_TEXT)

    @require_flash_attn
    @require_torch_accelerator
    @pytest.mark.flash_attn_test
    @slow
    def test_FA2(self):
        url = "https://huggingface.co/microsoft/kosmos-2.5/resolve/main/receipt_00008.png"
        image = Image.open(requests.get(url, stream=True).raw)

        dtype = torch.bfloat16
        repo = "microsoft/kosmos-2.5"
        model = Kosmos2_5ForConditionalGeneration.from_pretrained(
            repo,
            device_map=torch_device,
            dtype=dtype,
            attn_implementation="flash_attention_2",
        )
        processor = AutoProcessor.from_pretrained(repo)
        prompt = "<ocr>"
        generated_ids, generated_text = self.run_example(prompt, image, model, processor)
        EXPECTED_TEXT = [
            "<bbox><x_53><y_573><x_69><y_606></bbox>1\n<bbox><x_79><y_573><x_464><y_612></bbox>[REG] BLACK SAKURA\n<bbox><x_690><y_569><x_812><y_606></bbox>45,455\n<bbox><x_53><y_614><x_69><y_650></bbox>1\n<bbox><x_79><y_614><x_468><y_650></bbox>COOKIE DOH SAUCES\n<bbox><x_788><y_610><x_813><y_644></bbox>0\n<bbox><x_50><y_658><x_65><y_693></bbox>1\n<bbox><x_76><y_658><x_358><y_693></bbox>NATA DE COCO\n<bbox><x_790><y_652><x_815><y_687></bbox>0\n<bbox><x_31><y_742><x_822><y_781></bbox>Sub Total 45,455\n<bbox><x_27><y_780><x_822><y_827></bbox>PB1 (10%) 4,545\n<bbox><x_27><y_826><x_824><y_874></bbox>Rounding 0\n<bbox><x_24><y_872><x_827><y_921></bbox>Total 50,000\n<bbox><x_17><y_1056><x_835><y_1108></bbox>Card Payment 50,000\n"
        ]

        self.assertListEqual(generated_text, EXPECTED_TEXT)

        prompt = "<md>"
        generated_ids, generated_text = self.run_example(prompt, image, model, processor)
        # A10 gives the 1st one, but A100 gives the 2nd one
        EXPECTED_TEXT = [
            "- **1 \\[REG\\] BLACK SAKURA** 45,455\n- **1 COOKIE DOH SAUCES** 0\n- **1 NATA DE COCO** 0\n\n<table>\n<thead>\n<tr>\n<th>\nSub Total\n</th>\n<th>\n45,455\n</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>\nPB1 (10%)\n</td>\n<td>\n4,545\n</td>\n</tr>\n<tr>\n<td>\nRounding\n</td>\n<td>\n0\n</td>\n</tr>\n<tr>\n<td>\n<strong>\nTotal\n</strong>\n</td>\n<td>\n<strong>\n50,000\n</strong>\n</td>\n</tr>\n</tbody>\n</table>\n\nCard Payment 50,000",
            "- **1 \\[REG\\] BLACK SAKURA** 45,455\n- **1 COOKIE DOH SAUCES** 0\n- **1 NATA DE COCO** 0\n- **Sub Total** 45,455\n- **PB1 (10%)** 4,545\n- **Rounding** 0\n- **Total** **50,000**\n",
        ]
        self.assertIn(generated_text[0], EXPECTED_TEXT)