From 2b4694319539a006fdd89a85b0d91b4960b2e1ef Mon Sep 17 00:00:00 2001 From: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Date: Fri, 31 Jan 2025 11:28:13 -0500 Subject: [PATCH] Add GOT-OCR 2.0 to Transformers (#34721) * init modular got_ocr2 * Get correct got_ocr architecture * add processing * run modular with processing * add working inference * apply modular * Refactor and fix style * Refactor, cleanup, fix style * fix init order * Fix docs * add base modeling tests * fix style and consistency * rename doc file * fix repo consistency * fix inference with box * add image processing and support for crop_to_multi_page * Fix batch inference * add tests * fixup * fix slow test * fix docstrings * Add model doc * update to new init * fix input autocast pixel_values dtype * update doc * move doc to multimodal * Reformat crop_image_to_patches and add docstrings * Fix example in forward docstring * Address Pablo review * [run slow] got_ocr2 * remove defaults defined twice * apply modular * add torch_device to integration tests * update modular * follow-up Pavel review * add device variable in doc * fix doc multi-page * Force eager attention for vision encoder to avoid attn implementation conflict * revert qwen2vl doc changes * use Qwen2ForCausalLM instead of Qwen2Model * make fixup * refactor gotocr2 to llava style * uniformize function names and reduce checks * final nits * fix pixel_values dtype error * change checkpoint names * fix modular --- docs/source/en/_toctree.yml | 2 + docs/source/en/index.md | 1 + docs/source/en/model_doc/got_ocr2.md | 269 +++++ docs/source/en/perf_infer_gpu_one.md | 2 + src/transformers/__init__.py | 18 + src/transformers/models/__init__.py | 1 + .../models/auto/configuration_auto.py | 2 + .../models/auto/image_processing_auto.py | 1 + src/transformers/models/auto/modeling_auto.py | 3 + .../models/auto/processing_auto.py | 1 + src/transformers/models/got_ocr2/__init__.py | 31 + .../models/got_ocr2/configuration_got_ocr2.py | 212 ++++ .../convert_got_ocr2_weights_to_hf.py | 274 +++++ .../got_ocr2/image_processing_got_ocr2.py | 482 ++++++++ .../models/got_ocr2/modeling_got_ocr2.py | 1005 +++++++++++++++++ .../models/got_ocr2/modular_got_ocr2.py | 984 ++++++++++++++++ .../models/got_ocr2/processing_got_ocr2.py | 294 +++++ .../models/llava/modeling_llava.py | 2 +- src/transformers/utils/dummy_pt_objects.py | 14 + .../utils/dummy_vision_objects.py | 7 + tests/generation/test_utils.py | 2 +- tests/models/got_ocr2/__init__.py | 0 .../test_image_processing_got_ocr2.py | 115 ++ .../models/got_ocr2/test_modeling_got_ocr2.py | 386 +++++++ .../got_ocr2/test_processor_got_ocr2.py | 77 ++ .../models/qwen2_vl/test_modeling_qwen2_vl.py | 2 +- 26 files changed, 4184 insertions(+), 3 deletions(-) create mode 100644 docs/source/en/model_doc/got_ocr2.md create mode 100644 src/transformers/models/got_ocr2/__init__.py create mode 100644 src/transformers/models/got_ocr2/configuration_got_ocr2.py create mode 100644 src/transformers/models/got_ocr2/convert_got_ocr2_weights_to_hf.py create mode 100644 src/transformers/models/got_ocr2/image_processing_got_ocr2.py create mode 100644 src/transformers/models/got_ocr2/modeling_got_ocr2.py create mode 100644 src/transformers/models/got_ocr2/modular_got_ocr2.py create mode 100644 src/transformers/models/got_ocr2/processing_got_ocr2.py create mode 100644 tests/models/got_ocr2/__init__.py create mode 100644 tests/models/got_ocr2/test_image_processing_got_ocr2.py create mode 100644 tests/models/got_ocr2/test_modeling_got_ocr2.py create mode 100644 tests/models/got_ocr2/test_processor_got_ocr2.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 1e02abd1bb..2a2cf4512a 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -872,6 +872,8 @@ title: FLAVA - local: model_doc/git title: GIT + - local: model_doc/got_ocr2 + title: GOT-OCR2 - local: model_doc/grounding-dino title: Grounding DINO - local: model_doc/groupvit diff --git a/docs/source/en/index.md b/docs/source/en/index.md index 2233630128..92a082f0a9 100644 --- a/docs/source/en/index.md +++ b/docs/source/en/index.md @@ -161,6 +161,7 @@ Flax), PyTorch, and/or TensorFlow. | [GIT](model_doc/git) | ✅ | ❌ | ❌ | | [GLM](model_doc/glm) | ✅ | ❌ | ❌ | | [GLPN](model_doc/glpn) | ✅ | ❌ | ❌ | +| [GOT-OCR2](model_doc/got_ocr2) | ✅ | ❌ | ❌ | | [GPT Neo](model_doc/gpt_neo) | ✅ | ❌ | ✅ | | [GPT NeoX](model_doc/gpt_neox) | ✅ | ❌ | ❌ | | [GPT NeoX Japanese](model_doc/gpt_neox_japanese) | ✅ | ❌ | ❌ | diff --git a/docs/source/en/model_doc/got_ocr2.md b/docs/source/en/model_doc/got_ocr2.md new file mode 100644 index 0000000000..a560f78269 --- /dev/null +++ b/docs/source/en/model_doc/got_ocr2.md @@ -0,0 +1,269 @@ + + +# GOT-OCR2 + +## Overview + +The GOT-OCR2 model was proposed in [General OCR Theory: Towards OCR-2.0 via a Unified End-to-end Model](https://arxiv.org/abs/2409.01704) by Haoran Wei, Chenglong Liu, Jinyue Chen, Jia Wang, Lingyu Kong, Yanming Xu, Zheng Ge, Liang Zhao, Jianjian Sun, Yuang Peng, Chunrui Han, Xiangyu Zhang. + +The abstract from the paper is the following: + +*Traditional OCR systems (OCR-1.0) are increasingly unable to meet people’snusage due to the growing demand for intelligent processing of man-made opticalncharacters. In this paper, we collectively refer to all artificial optical signals (e.g., plain texts, math/molecular formulas, tables, charts, sheet music, and even geometric shapes) as "characters" and propose the General OCR Theory along with an excellent model, namely GOT, to promote the arrival of OCR-2.0. The GOT, with 580M parameters, is a unified, elegant, and end-to-end model, consisting of a high-compression encoder and a long-contexts decoder. As an OCR-2.0 model, GOT can handle all the above "characters" under various OCR tasks. On the input side, the model supports commonly used scene- and document-style images in slice and whole-page styles. On the output side, GOT can generate plain or formatted results (markdown/tikz/smiles/kern) via an easy prompt. Besides, the model enjoys interactive OCR features, i.e., region-level recognition guided by coordinates or colors. Furthermore, we also adapt dynamic resolution and multipage OCR technologies to GOT for better practicality. In experiments, we provide sufficient results to prove the superiority of our model.* + + + + GOT-OCR2 training stages. Taken from the original paper. + + +Tips: + +GOT-OCR2 works on a wide range of tasks, including plain document OCR, scene text OCR, formatted document OCR, and even OCR for tables, charts, mathematical formulas, geometric shapes, molecular formulas and sheet music. While this implementation of the model will only output plain text, the outputs can be further processed to render the desired format, with packages like `pdftex`, `mathpix`, `matplotlib`, `tikz`, `verovio` or `pyecharts`. +The model can also be used for interactive OCR, where the user can specify the region to be recognized by providing the coordinates or the color of the region's bounding box. + +This model was contributed by [yonigozlan](https://huggingface.co/yonigozlan). +The original code can be found [here](https://github.com/Ucas-HaoranWei/GOT-OCR2.0). + +## Usage example + +### Plain text inference + +```python +>>> from transformers import AutoProcessor, AutoModelForImageTextToText + +>>> device = "cuda" if torch.cuda.is_available() else "cpu" +>>> model = AutoModelForImageTextToText.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf", device_map=device) +>>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + +>>> image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" +>>> inputs = processor(image, return_tensors="pt").to(device) + +>>> generate_ids = model.generate( +... **inputs, +... do_sample=False, +... tokenizer=processor.tokenizer, +... stop_strings="<|im_end|>", +... max_new_tokens=4096, +... ) + +>>> processor.decode(generate_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +"R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM\nName/Phone Ext. : (...)" +``` + +### Plain text inference batched + +```python +>>> from transformers import AutoProcessor, AutoModelForImageTextToText + +>>> device = "cuda" if torch.cuda.is_available() else "cpu" +>>> model = AutoModelForImageTextToText.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf", device_map=device) +>>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + +>>> image1 = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" +>>> image2 = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" + +>>> inputs = processor([image1, image2], return_tensors="pt").to(device) + +>>> generate_ids = model.generate( +... **inputs, +... do_sample=False, +... tokenizer=processor.tokenizer, +... stop_strings="<|im_end|>", +... max_new_tokens=4, +... ) + +>>> processor.batch_decode(generate_ids[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True) +["Reducing the number", "R&D QUALITY"] +``` + +### Formatted text inference + +GOT-OCR2 can also generate formatted text, such as markdown or LaTeX. Here is an example of how to generate formatted text: + +```python +>>> from transformers import AutoProcessor, AutoModelForImageTextToText + +>>> device = "cuda" if torch.cuda.is_available() else "cpu" +>>> model = AutoModelForImageTextToText.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf", device_map=device) +>>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + +>>> image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/latex.png" +>>> inputs = processor(image, return_tensors="pt", format=True).to(device) + +>>> generate_ids = model.generate( +... **inputs, +... do_sample=False, +... tokenizer=processor.tokenizer, +... stop_strings="<|im_end|>", +... max_new_tokens=4096, +... ) + +>>> processor.decode(generate_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +"\\author{\nHanwen Jiang* \\(\\quad\\) Arjun Karpur \\({ }^{\\dagger} \\quad\\) Bingyi Cao \\({ }^{\\dagger} \\quad\\) (...)" +``` + +### Inference on multiple pages + +Although it might be reasonable in most cases to use a “for loop” for multi-page processing, some text data with formatting across several pages make it necessary to process all pages at once. GOT introduces a multi-page OCR (without “for loop”) feature, where multiple pages can be processed by the model at once, whith the output being one continuous text. +Here is an example of how to process multiple pages at once: + + +```python +>>> from transformers import AutoProcessor, AutoModelForImageTextToText + +>>> device = "cuda" if torch.cuda.is_available() else "cpu" +>>> model = AutoModelForImageTextToText.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf", device_map=device) +>>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + +>>> image1 = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/page1.png" +>>> image2 = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/page2.png" +>>> inputs = processor([image1, image2], return_tensors="pt", multi_page=True, format=True).to(device) + +>>> generate_ids = model.generate( +... **inputs, +... do_sample=False, +... tokenizer=processor.tokenizer, +... stop_strings="<|im_end|>", +... max_new_tokens=4096, +... ) + +>>> processor.decode(generate_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +"\\title{\nGeneral OCR Theory: Towards OCR-2.0 via a Unified End-to-end Model\n}\n\\author{\nHaoran Wei (...)" +``` + +### Inference on cropped patches + +GOT supports a 1024×1024 input resolution, which is sufficient for most OCR tasks, such as scene OCR or processing A4-sized PDF pages. However, certain scenarios, like horizontally stitched two-page PDFs commonly found in academic papers or images with unusual aspect ratios, can lead to accuracy issues when processed as a single image. To address this, GOT can dynamically crop an image into patches, process them all at once, and merge the results for better accuracy with such inputs. +Here is an example of how to process cropped patches: + +```python +>>> import torch +>>> from transformers import AutoProcessor, AutoModelForImageTextToText + +>>> device = "cuda" if torch.cuda.is_available() else "cpu" +>>> model = AutoModelForImageTextToText.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf", torch_dtype=torch.bfloat16, device_map=device) +>>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + +>>> image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/one_column.png" +>>> inputs = processor(image, return_tensors="pt", format=True, crop_to_patches=True, max_patches=3).to(device) + +>>> generate_ids = model.generate( +... **inputs, +... do_sample=False, +... tokenizer=processor.tokenizer, +... stop_strings="<|im_end|>", +... max_new_tokens=4096, +... ) + +>>> processor.decode(generate_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +"on developing architectural improvements to make learnable matching methods generalize.\nMotivated by the above observations, (...)" +``` + +### Inference on a specific region + +GOT supports interactive OCR, where the user can specify the region to be recognized by providing the coordinates or the color of the region's bounding box. Here is an example of how to process a specific region: + +```python +>>> from transformers import AutoProcessor, AutoModelForImageTextToText + +>>> device = "cuda" if torch.cuda.is_available() else "cpu" +>>> model = AutoModelForImageTextToText.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf", device_map=device) +>>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + +>>> image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" +>>> inputs = processor(image, return_tensors="pt", color="green").to(device) # or box=[x1, y1, x2, y2] for coordinates (image pixels) + +>>> generate_ids = model.generate( +... **inputs, +... do_sample=False, +... tokenizer=processor.tokenizer, +... stop_strings="<|im_end|>", +... max_new_tokens=4096, +... ) + +>>> processor.decode(generate_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +"You should keep in mind what features from the module should be used, especially \nwhen you’re planning to sell a template." +``` + +### Inference on general OCR data example: sheet music + +Although this implementation of the model will only output plain text, the outputs can be further processed to render the desired format, with packages like `pdftex`, `mathpix`, `matplotlib`, `tikz`, `verovio` or `pyecharts`. +Here is an example of how to process sheet music: + +```python +>>> from transformers import AutoProcessor, AutoModelForImageTextToText +>>> import verovio + +>>> device = "cuda" if torch.cuda.is_available() else "cpu" +>>> model = AutoModelForImageTextToText.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf", device_map=device) +>>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + +>>> image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/sheet_music.png" +>>> inputs = processor(image, return_tensors="pt", format=True).to(device) + +>>> generate_ids = model.generate( +... **inputs, +... do_sample=False, +... tokenizer=processor.tokenizer, +... stop_strings="<|im_end|>", +... max_new_tokens=4096, +... ) + +>>> outputs = processor.decode(generate_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +>>> tk = verovio.toolkit() +>>> tk.loadData(outputs) +>>> tk.setOptions( +... { +... "pageWidth": 2100, +... "pageHeight": 800, +... "footer": "none", +... "barLineWidth": 0.5, +... "beamMaxSlope": 15, +... "staffLineWidth": 0.2, +... "spacingStaff": 6, +... } +... ) +>>> tk.getPageCount() +>>> svg = tk.renderToSVG() +>>> svg = svg.replace('overflow="inherit"', 'overflow="visible"') +>>> with open("output.svg", "w") as f: +>>> f.write(svg) +``` + + +## GotOcr2Config + +[[autodoc]] GotOcr2Config + +## GotOcr2VisionConfig + +[[autodoc]] GotOcr2VisionConfig + +## GotOcr2ImageProcessor + +[[autodoc]] GotOcr2ImageProcessor + +## GotOcr2Processor + +[[autodoc]] GotOcr2Processor + +## GotOcr2ForConditionalGeneration + +[[autodoc]] GotOcr2ForConditionalGeneration + - forward + diff --git a/docs/source/en/perf_infer_gpu_one.md b/docs/source/en/perf_infer_gpu_one.md index 8087008f87..5ffdefec9e 100644 --- a/docs/source/en/perf_infer_gpu_one.md +++ b/docs/source/en/perf_infer_gpu_one.md @@ -52,6 +52,7 @@ FlashAttention-2 is currently supported for the following architectures: * [Emu3](https://huggingface.co/docs/transformers/model_doc/emu3) * [Gemma](https://huggingface.co/docs/transformers/model_doc/gemma#transformers.GemmaModel) * [Gemma2](https://huggingface.co/docs/transformers/model_doc/gemma2#transformers.Gemma2Model) +* [GotOcr2](https://huggingface.co/docs/transformers/model_doc/got_ocr2#transformers.GotOcr2ForConditionalGeneration) * [GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2) * [GPTBigCode](https://huggingface.co/docs/transformers/model_doc/gpt_bigcode#transformers.GPTBigCodeModel) * [GPTNeo](https://huggingface.co/docs/transformers/model_doc/gpt_neo#transformers.GPTNeoModel) @@ -253,6 +254,7 @@ For now, Transformers supports SDPA inference and training for the following arc * [Falcon](https://huggingface.co/docs/transformers/model_doc/falcon#transformers.FalconModel) * [Gemma](https://huggingface.co/docs/transformers/model_doc/gemma#transformers.GemmaModel) * [Gemma2](https://huggingface.co/docs/transformers/model_doc/gemma2#transformers.Gemma2Model) +* [GotOcr2](https://huggingface.co/docs/transformers/model_doc/got_ocr2#transformers.GotOcr2ForConditionalGeneration) * [Granite](https://huggingface.co/docs/transformers/model_doc/granite#transformers.GraniteModel) * [GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2) * [GPTBigCode](https://huggingface.co/docs/transformers/model_doc/gpt_bigcode#transformers.GPTBigCodeModel) diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 02a55c28ac..ae92f21dcc 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -476,6 +476,11 @@ _import_structure = { ], "models.glm": ["GlmConfig"], "models.glpn": ["GLPNConfig"], + "models.got_ocr2": [ + "GotOcr2Config", + "GotOcr2Processor", + "GotOcr2VisionConfig", + ], "models.gpt2": [ "GPT2Config", "GPT2Tokenizer", @@ -1238,6 +1243,7 @@ else: _import_structure["models.flava"].extend(["FlavaFeatureExtractor", "FlavaImageProcessor", "FlavaProcessor"]) _import_structure["models.fuyu"].extend(["FuyuImageProcessor", "FuyuProcessor"]) _import_structure["models.glpn"].extend(["GLPNFeatureExtractor", "GLPNImageProcessor"]) + _import_structure["models.got_ocr2"].extend(["GotOcr2ImageProcessor"]) _import_structure["models.grounding_dino"].extend(["GroundingDinoImageProcessor"]) _import_structure["models.idefics"].extend(["IdeficsImageProcessor"]) _import_structure["models.idefics2"].extend(["Idefics2ImageProcessor"]) @@ -2426,6 +2432,12 @@ else: "GLPNPreTrainedModel", ] ) + _import_structure["models.got_ocr2"].extend( + [ + "GotOcr2ForConditionalGeneration", + "GotOcr2PreTrainedModel", + ] + ) _import_structure["models.gpt2"].extend( [ "GPT2DoubleHeadsModel", @@ -5540,6 +5552,7 @@ if TYPE_CHECKING: ) from .models.glm import GlmConfig from .models.glpn import GLPNConfig + from .models.got_ocr2 import GotOcr2Config, GotOcr2Processor, GotOcr2VisionConfig from .models.gpt2 import ( GPT2Config, GPT2Tokenizer, @@ -6342,6 +6355,7 @@ if TYPE_CHECKING: ) from .models.fuyu import FuyuImageProcessor, FuyuProcessor from .models.glpn import GLPNFeatureExtractor, GLPNImageProcessor + from .models.got_ocr2 import GotOcr2ImageProcessor from .models.grounding_dino import GroundingDinoImageProcessor from .models.idefics import IdeficsImageProcessor from .models.idefics2 import Idefics2ImageProcessor @@ -7346,6 +7360,10 @@ if TYPE_CHECKING: GLPNModel, GLPNPreTrainedModel, ) + from .models.got_ocr2 import ( + GotOcr2ForConditionalGeneration, + GotOcr2PreTrainedModel, + ) from .models.gpt2 import ( GPT2DoubleHeadsModel, GPT2ForQuestionAnswering, diff --git a/src/transformers/models/__init__.py b/src/transformers/models/__init__.py index 43cf2fe429..f62d5d7167 100644 --- a/src/transformers/models/__init__.py +++ b/src/transformers/models/__init__.py @@ -106,6 +106,7 @@ from . import ( git, glm, glpn, + got_ocr2, gpt2, gpt_bigcode, gpt_neo, diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py index 95c5bc4d00..699e307ac1 100644 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py @@ -124,6 +124,7 @@ CONFIG_MAPPING_NAMES = OrderedDict( ("git", "GitConfig"), ("glm", "GlmConfig"), ("glpn", "GLPNConfig"), + ("got_ocr2", "GotOcr2Config"), ("gpt-sw3", "GPT2Config"), ("gpt2", "GPT2Config"), ("gpt_bigcode", "GPTBigCodeConfig"), @@ -450,6 +451,7 @@ MODEL_NAMES_MAPPING = OrderedDict( ("git", "GIT"), ("glm", "GLM"), ("glpn", "GLPN"), + ("got_ocr2", "GOT-OCR2"), ("gpt-sw3", "GPT-Sw3"), ("gpt2", "OpenAI GPT-2"), ("gpt_bigcode", "GPTBigCode"), diff --git a/src/transformers/models/auto/image_processing_auto.py b/src/transformers/models/auto/image_processing_auto.py index 1cb067f386..0c4bb9b9a7 100644 --- a/src/transformers/models/auto/image_processing_auto.py +++ b/src/transformers/models/auto/image_processing_auto.py @@ -87,6 +87,7 @@ else: ("fuyu", ("FuyuImageProcessor",)), ("git", ("CLIPImageProcessor",)), ("glpn", ("GLPNImageProcessor",)), + ("got_ocr2", ("GotOcr2ImageProcessor",)), ("grounding-dino", ("GroundingDinoImageProcessor",)), ("groupvit", ("CLIPImageProcessor",)), ("hiera", ("BitImageProcessor",)), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 8030f5dbbd..3b023251e1 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -119,6 +119,7 @@ MODEL_MAPPING_NAMES = OrderedDict( ("git", "GitModel"), ("glm", "GlmModel"), ("glpn", "GLPNModel"), + ("got_ocr2", "GotOcr2ForConditionalGeneration"), ("gpt-sw3", "GPT2Model"), ("gpt2", "GPT2Model"), ("gpt_bigcode", "GPTBigCodeModel"), @@ -512,6 +513,7 @@ MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict( ("gemma2", "Gemma2ForCausalLM"), ("git", "GitForCausalLM"), ("glm", "GlmForCausalLM"), + ("got_ocr2", "GotOcr2ForConditionalGeneration"), ("gpt-sw3", "GPT2LMHeadModel"), ("gpt2", "GPT2LMHeadModel"), ("gpt_bigcode", "GPTBigCodeForCausalLM"), @@ -811,6 +813,7 @@ MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES = OrderedDict( ("emu3", "Emu3ForConditionalGeneration"), ("fuyu", "FuyuForCausalLM"), ("git", "GitForCausalLM"), + ("got_ocr2", "GotOcr2ForConditionalGeneration"), ("idefics", "IdeficsForVisionText2Text"), ("idefics2", "Idefics2ForConditionalGeneration"), ("idefics3", "Idefics3ForConditionalGeneration"), diff --git a/src/transformers/models/auto/processing_auto.py b/src/transformers/models/auto/processing_auto.py index 678b66a397..f329d9e465 100644 --- a/src/transformers/models/auto/processing_auto.py +++ b/src/transformers/models/auto/processing_auto.py @@ -63,6 +63,7 @@ PROCESSOR_MAPPING_NAMES = OrderedDict( ("flava", "FlavaProcessor"), ("fuyu", "FuyuProcessor"), ("git", "GitProcessor"), + ("got_ocr2", "GotOcr2Processor"), ("grounding-dino", "GroundingDinoProcessor"), ("groupvit", "CLIPProcessor"), ("hubert", "Wav2Vec2Processor"), diff --git a/src/transformers/models/got_ocr2/__init__.py b/src/transformers/models/got_ocr2/__init__.py new file mode 100644 index 0000000000..071a7ea740 --- /dev/null +++ b/src/transformers/models/got_ocr2/__init__.py @@ -0,0 +1,31 @@ +# Copyright 2024 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. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_got_ocr2 import * + from .image_processing_got_ocr2 import * + from .modeling_got_ocr2 import * + from .processing_got_ocr2 import * + + +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/src/transformers/models/got_ocr2/configuration_got_ocr2.py b/src/transformers/models/got_ocr2/configuration_got_ocr2.py new file mode 100644 index 0000000000..480252ab14 --- /dev/null +++ b/src/transformers/models/got_ocr2/configuration_got_ocr2.py @@ -0,0 +1,212 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/got_ocr2/modular_got_ocr2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_got_ocr2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2024 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. + + +from ...configuration_utils import PretrainedConfig +from ..auto import CONFIG_MAPPING, AutoConfig + + +class GotOcr2VisionConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`GotOcr2VisionModel`]. It is used to instantiate a GOT_OCR2 + vision encoder according to the specified arguments, defining the model architecture. Instantiating a configuration + defaults will yield a similar configuration to that of the SAM ViT-h + [facebook/sam-vit-huge](https://huggingface.co/facebook/sam-vit-huge) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + hidden_size (`int`, *optional*, defaults to 768): + Dimensionality of the encoder layers and the pooler layer. + output_channels (`int`, *optional*, defaults to 256): + Dimensionality of the output channels in the Patch Encoder. + num_hidden_layers (`int`, *optional*, defaults to 12): + Number of hidden layers in the Transformer encoder. + num_attention_heads (`int`, *optional*, defaults to 12): + Number of attention heads for each attention layer in the Transformer encoder. + num_channels (`int`, *optional*, defaults to 3): + Number of channels in the input image. + image_size (`int`, *optional*, defaults to 1024): + Expected resolution. Target size of the resized input image. + patch_size (`int`, *optional*, defaults to 16): + Size of the patches to be extracted from the input image. + hidden_act (`str`, *optional*, defaults to `"gelu"`): + The non-linear activation function (function or string) + layer_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the layer normalization layers. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + initializer_range (`float`, *optional*, defaults to 1e-10): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + qkv_bias (`bool`, *optional*, defaults to `True`): + Whether to add a bias to query, key, value projections. + use_abs_pos (`bool`, *optional*, defaults to `True`): + Whether to use absolute position embedding. + use_rel_pos (`bool`, *optional*, defaults to `True`): + Whether to use relative position embedding. + window_size (`int`, *optional*, defaults to 14): + Window size for relative position. + global_attn_indexes (`List[int]`, *optional*, defaults to `[2, 5, 8, 11]`): + The indexes of the global attention layers. + mlp_dim (`int`, *optional*, defaults to 3072): + The dimensionality of the MLP layer in the Transformer encoder. + """ + + base_config_key = "vision_config" + + def __init__( + self, + hidden_size=768, + output_channels=256, + num_hidden_layers=12, + num_attention_heads=12, + num_channels=3, + image_size=1024, + patch_size=16, + hidden_act="gelu", + layer_norm_eps=1e-06, + attention_dropout=0.0, + initializer_range=1e-10, + qkv_bias=True, + use_abs_pos=True, + use_rel_pos=True, + window_size=14, + global_attn_indexes=[2, 5, 8, 11], + mlp_dim=3072, + **kwargs, + ): + super().__init__(**kwargs) + + self.hidden_size = hidden_size + self.output_channels = output_channels + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.hidden_act = hidden_act + self.layer_norm_eps = layer_norm_eps + self.attention_dropout = attention_dropout + self.initializer_range = initializer_range + self.qkv_bias = qkv_bias + self.use_abs_pos = use_abs_pos + self.use_rel_pos = use_rel_pos + self.window_size = window_size + self.global_attn_indexes = global_attn_indexes + self.mlp_dim = mlp_dim + + +class GotOcr2Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`GotOcr2ForConditionalGeneration`]. It is used to instantiate a + GotOcr2 model according to the specified arguments, defining the model architecture. Instantiating a configuration + with the defaults will yield a similar configuration to that of GOT-OCR-2.0. + + e.g [stepfun-ai/GOT-OCR-2.0-hf](https://huggingface.co/stepfun-ai/GOT-OCR-2.0-hf) + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vision_config (`Union[AutoConfig, dict]`, *optional*, defaults to `CLIPVisionConfig`): + The config object or dictionary of the vision backbone. + text_config (`Union[AutoConfig, dict]`, *optional*, defaults to `LlamaConfig`): + The config object or dictionary of the text backbone. + ignore_index (`int`, *optional*, defaults to -100): + The ignore index for the loss function. + image_token_index (`int`, *optional*, defaults to 151859): + The image token index to encode the image prompt. + image_seq_length (`int`, *optional*, defaults to 576): + Sequence length of one image embedding. + pad_token_id (`int`, *optional*, defaults to -1): + Padding token id. + + ```python + >>> from transformers import GotOcr2ForConditionalGeneration, GotOcr2Config + + >>> # Initializing a GotOcr2 style configuration + >>> configuration = GotOcr2Config() + + >>> # Initializing a model from the Qwen2-VL-7B style configuration + >>> model = GotOcr2ForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "got_ocr2" + sub_configs = {"text_config": AutoConfig, "vision_config": GotOcr2VisionConfig} + + def __init__( + self, + vision_config=None, + text_config=None, + ignore_index=-100, + image_token_index=151859, + image_seq_length=576, + pad_token_id=-1, + **kwargs, + ): + self.ignore_index = ignore_index + self.image_token_index = image_token_index + self.image_seq_length = image_seq_length + self.pad_token_id = pad_token_id + + if vision_config is None: + self.vision_config = GotOcr2VisionConfig() + elif isinstance(vision_config, dict): + self.vision_config = GotOcr2VisionConfig(**vision_config) + elif isinstance(vision_config, GotOcr2VisionConfig): + self.vision_config = vision_config + + if isinstance(text_config, dict): + text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "qwen2" + text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config) + elif text_config is None: + text_config = CONFIG_MAPPING["qwen2"]( + vocab_size=151860, + hidden_size=1024, + intermediate_size=2816, + num_hidden_layers=24, + num_attention_heads=16, + num_key_value_heads=16, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=True, + rope_theta=1000000.0, + rope_scaling=None, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=21, + attention_dropout=0.0, + ) + + self.text_config = text_config + + super().__init__(**kwargs) + + +__all__ = ["GotOcr2VisionConfig", "GotOcr2Config"] diff --git a/src/transformers/models/got_ocr2/convert_got_ocr2_weights_to_hf.py b/src/transformers/models/got_ocr2/convert_got_ocr2_weights_to_hf.py new file mode 100644 index 0000000000..3df7214410 --- /dev/null +++ b/src/transformers/models/got_ocr2/convert_got_ocr2_weights_to_hf.py @@ -0,0 +1,274 @@ +# Copyright 2024 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. + +import argparse +import gc +import glob +import os +from typing import List, Optional + +import regex as re +import torch +from huggingface_hub import snapshot_download +from safetensors import safe_open + +from transformers import ( + GotOcr2Config, + GotOcr2ForConditionalGeneration, + GotOcr2ImageProcessor, + GotOcr2Processor, + PreTrainedTokenizerFast, + is_vision_available, +) +from transformers.convert_slow_tokenizer import TikTokenConverter +from transformers.tokenization_utils import AddedToken + + +if is_vision_available(): + from transformers.image_utils import load_image + + +# fmt: off +ORIGINAL_TO_CONVERTED_KEY_MAPPING = { + # Vision encoder mapping + r"model.vision_tower_high.pos_embed": r"vision_tower.pos_embed", + r"model.vision_tower_high.patch_embed.proj": r"vision_tower.patch_embed.projection", + r"model.vision_tower_high.blocks.(\d+).norm": r"vision_tower.layers.\1.layer_norm", + r"model.vision_tower_high.blocks.(\d+).attn": r"vision_tower.layers.\1.attn", + r"model.vision_tower_high.blocks.(\d+).mlp": r"vision_tower.layers.\1.mlp", + r"model.vision_tower_high.neck.0": r"vision_tower.neck.conv1", + r"model.vision_tower_high.neck.1": r"vision_tower.neck.layer_norm1", + r"model.vision_tower_high.neck.2": r"vision_tower.neck.conv2", + r"model.vision_tower_high.neck.3": r"vision_tower.neck.layer_norm2", + r"model.vision_tower_high.net_(\d+)": lambda m: f"multi_modal_projector.conv_upsampler{int(m.group(1)) - 1}", + r"model.mm_projector_vary" : r"multi_modal_projector.multimodal_projector", + r"model.": r"language_model.model.", + r"lm_head": r"language_model.lm_head", +} +# fmt: on + +CONTEXT_LENGTH = 8000 + + +def convert_old_keys_to_new_keys(state_dict_keys: dict = None): + """ + This function should be applied only once, on the concatenated keys to efficiently rename using + the key mappings. + """ + output_dict = {} + if state_dict_keys is not None: + old_text = "\n".join(state_dict_keys) + new_text = old_text + for pattern, replacement in ORIGINAL_TO_CONVERTED_KEY_MAPPING.items(): + new_text = re.sub(pattern, replacement, new_text) + output_dict = dict(zip(old_text.split("\n"), new_text.split("\n"))) + return output_dict + + +def load_original_state_dict(model_id): + directory_path = snapshot_download(repo_id=model_id, allow_patterns=["*.safetensors"]) + + original_state_dict = {} + for path in glob.glob(f"{directory_path}/*"): + if path.endswith(".safetensors"): + with safe_open(path, framework="pt", device="cpu") as f: + for key in f.keys(): + original_state_dict[key] = f.get_tensor(key) + + return original_state_dict + + +def get_got_ocr2_config(): + config = GotOcr2Config() + + return config + + +def write_model( + model_path, + input_base_path, + push_to_hub=False, +): + os.makedirs(model_path, exist_ok=True) + + config = get_got_ocr2_config() + config.architectures = ["GotOcr2ForConditionalGeneration"] + config.save_pretrained(model_path) + print("Model config saved successfully...") + + # ------------------------------------------------------------ + # Convert weights + # ------------------------------------------------------------ + + print(f"Fetching all parameters from the checkpoint at {input_base_path}...") + state_dict_old = load_original_state_dict(input_base_path) + print("Converting model...") + all_keys = list(state_dict_old.keys()) + new_keys = convert_old_keys_to_new_keys(all_keys) + state_dict = {} + for key in all_keys: + new_key = new_keys[key] + state_dict[new_key] = state_dict_old[key] + + del state_dict_old + gc.collect() + + print("Loading the checkpoint in a GotOcr2ForConditionalGeneration model.") + model = GotOcr2ForConditionalGeneration(config) + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + model = model.to(torch.bfloat16) + print("model dtype:", model.dtype) + print("Missing keys:", missing_keys) + print("Unexpected keys:", unexpected_keys) + + print("Saving the model.") + model.save_pretrained(model_path) + if push_to_hub: + model.push_to_hub("stepfun-ai/GOT-OCR-2.0-hf", use_temp_dir=True) + del state_dict, model + + # Safety check: reload the converted model + gc.collect() + print("Reloading the model to check if it's saved correctly.") + model = GotOcr2ForConditionalGeneration.from_pretrained(model_path, device_map="auto") + processor = GotOcr2Processor.from_pretrained(model_path) + image = load_image( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" + ) + + inputs = processor(image, return_tensors="pt", format=True).to(model.device, dtype=model.dtype) + generate_ids = model.generate(**inputs, do_sample=False, num_beams=1, max_new_tokens=4) + decoded_output = processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) + expected_output = "\\title{\nR" + print("Decoded output:", decoded_output) + assert decoded_output == expected_output + print("Model reloaded successfully.") + del model + + +class GotOcr2Converter(TikTokenConverter): + def __init__( + self, + vocab_file, + special_tokens: List[str], + pattern: str, + model_max_length: int, + chat_template: Optional[str] = None, + **kwargs, + ): + super().__init__(vocab_file, pattern=pattern) + self.additional_special_tokens = special_tokens + tokenizer = self.converted() + if chat_template is not None: + kwargs["chat_template"] = chat_template + self.tokenizer = PreTrainedTokenizerFast( + tokenizer_object=tokenizer, + model_input_names=["input_ids", "attention_mask"], + model_max_length=model_max_length, + **kwargs, + ) + + +def write_tokenizer(tokenizer_path: str, save_dir: str, push_to_hub: bool = False): + model_max_length = CONTEXT_LENGTH + pattern = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" # noqa: W605 + # Special tokens + special_tokens = ( + ["<|endoftext|>", "<|im_start|>", "<|im_end|>"] + + [f"<|extra_{i}|>" for i in range(205)] + + [ + "", + "", + "", + "", + "", + "", + "", + "", + "", + ] + ) + + pad_token = "<|endoftext|>" + pad_token = AddedToken(pad_token, lstrip=False, rstrip=False, normalized=False, single_word=False) + + converter = GotOcr2Converter( + vocab_file=tokenizer_path, + pattern=pattern, + special_tokens=special_tokens, + model_max_length=model_max_length, + pad_token=pad_token, + bos_token="<|endoftext|>", + eos_token="<|endoftext|>", + clean_up_tokenization_spaces=True, + ) + tokenizer = converter.tokenizer + tokenizer.save_pretrained(save_dir) + + if push_to_hub: + tokenizer.push_to_hub("stepfun-ai/GOT-OCR-2.0-hf", use_temp_dir=True) + + +def write_image_processor(save_dir: str, push_to_hub: bool = False): + image_processor = GotOcr2ImageProcessor( + do_resize=True, + size={"height": 1024, "width": 1024}, + do_rescale=True, + rescale_factor=1 / 255, + do_normalize=True, + image_mean=[0.48145466, 0.4578275, 0.40821073], + image_std=[0.26862954, 0.26130258, 0.27577711], + ) + + image_processor.save_pretrained(save_dir) + if push_to_hub: + image_processor.push_to_hub("stepfun-ai/GOT-OCR-2.0-hf", use_temp_dir=True) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_dir", + default="stepfun-ai/GOT-OCR2_0", + help="Location of LLaMA weights, which contains tokenizer.model and model folders", + ) + parser.add_argument( + "--output_dir", + default="GotOcr2", + help="Location to write HF model and tokenizer", + ) + + parser.add_argument( + "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub." + ) + args = parser.parse_args() + write_tokenizer( + tokenizer_path="qwen.tiktoken", + save_dir=args.output_dir, + push_to_hub=args.push_to_hub, + ) + + write_image_processor( + save_dir=args.output_dir, + push_to_hub=args.push_to_hub, + ) + write_model( + model_path=args.output_dir, + input_base_path=args.input_dir, + push_to_hub=args.push_to_hub, + ) + + +if __name__ == "__main__": + main() diff --git a/src/transformers/models/got_ocr2/image_processing_got_ocr2.py b/src/transformers/models/got_ocr2/image_processing_got_ocr2.py new file mode 100644 index 0000000000..7f7a0d7ae4 --- /dev/null +++ b/src/transformers/models/got_ocr2/image_processing_got_ocr2.py @@ -0,0 +1,482 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/got_ocr2/modular_got_ocr2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_got_ocr2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2024 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. + + +from functools import lru_cache +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np + +from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict +from ...image_transforms import ( + _rescale_for_pil_conversion, + convert_to_rgb, + resize, + to_channel_dimension_format, + to_pil_image, +) +from ...image_utils import ( + OPENAI_CLIP_MEAN, + OPENAI_CLIP_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + infer_channel_dimension_format, + is_scaled_image, + make_flat_list_of_images, + to_numpy_array, + valid_images, + validate_preprocess_arguments, +) +from ...utils import TensorType, filter_out_non_signature_kwargs, is_vision_available, logging + + +if is_vision_available(): + import PIL + + +logger = logging.get_logger(__name__) + + +# Similar to image_processing_mllama.get_all_supported_aspect_ratios +@lru_cache(maxsize=10) +def get_all_supported_aspect_ratios(min_image_tiles: int, max_image_tiles: int) -> List[Tuple[int, int]]: + """ + Computes all allowed aspect ratios for a given minimum and maximum number of input tiles. + + This function calculates all possible arrangements of tiles that can be formed + within the constraint of the minimum and maximum number of tiles. Each arrangement is + represented by its aspect ratio (width/height) and the corresponding tile configuration. + + Args: + min_image_tiles (`int`): + The minimum number of tiles allowed. + max_image_tiles (`int`): + The maximum number of tiles allowed. + + Returns: + `List[Tuple[int, int]]`: A list of tuples, each tuple representing a valid (width, height) + configuration in terms of number of tiles. + + Example: + >>> get_all_supported_aspect_ratios(1, 4) + [(1, 1), (1, 2), (2, 1), (1, 3), (3, 1), (1, 4), (2, 2), (4, 1)] + + """ + aspect_ratios = [] + for width in range(1, max_image_tiles + 1): + for height in range(1, max_image_tiles + 1): + if width * height <= max_image_tiles and width * height >= min_image_tiles: + aspect_ratios.append((width, height)) + + aspect_ratios = sorted(aspect_ratios, key=lambda x: x[0] * x[1]) + + return aspect_ratios + + +@lru_cache(maxsize=100) +def get_optimal_tiled_canvas( + original_image_size: Tuple[int, int], + target_tile_size: Tuple[int, int], + min_image_tiles: int, + max_image_tiles: int, +) -> Tuple[int, int]: + """ + Given a minimum and maximum number of tiles, find the canvas with the closest aspect ratio to the + original image aspect ratio. + In case of tie-breaking condition when two canvases have the same aspect ratio difference, we favor the canvas with + more tiles, until the area covered by the tiles is more than twice the target area, in order to avoid unnecessarily + excessive tiling. + """ + possible_tile_arrangements = get_all_supported_aspect_ratios(min_image_tiles, max_image_tiles) + + original_height, original_width = original_image_size + target_tile_height, target_tile_width = target_tile_size + aspect_ratio = original_width / original_height + area = original_width * original_height + + # find the grid with the best aspect ratio + best_ratio_diff = float("inf") + best_grid = (1, 1) + for grid in possible_tile_arrangements: + grid_aspect_ratio = grid[0] / grid[1] + ratio_diff = abs(aspect_ratio - grid_aspect_ratio) + if ratio_diff < best_ratio_diff: + best_ratio_diff = ratio_diff + best_grid = grid + elif ratio_diff == best_ratio_diff: + # if the aspect ratio difference is the same, we favor the grid with more patches + # until the area covered by the patches is more than twice the original image area + if area > 0.5 * target_tile_height * target_tile_width * grid[0] * grid[1]: + best_grid = grid + + return best_grid + + +class GotOcr2ImageProcessor(BaseImageProcessor): + r""" + Constructs a GOT_OCR2 image processor. + + Args: + do_resize (`bool`, *optional*, defaults to `True`): + Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by the + `do_resize` parameter in the `preprocess` method. + size (`dict`, *optional*, defaults to `{"height": 384, "width": 384}`): + Size of the output image after resizing. Can be overridden by the `size` parameter in the `preprocess` + method. + resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`): + Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`. Can be + overridden by the `resample` parameter in the `preprocess` method. + do_rescale (`bool`, *optional*, defaults to `True`): + Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the + `do_rescale` parameter in the `preprocess` method. + rescale_factor (`int` or `float`, *optional*, defaults to `1/255`): + Scale factor to use if rescaling the image. Only has an effect if `do_rescale` is set to `True`. Can be + overridden by the `rescale_factor` parameter in the `preprocess` method. + do_normalize (`bool`, *optional*, defaults to `True`): + Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess` + method. Can be overridden by the `do_normalize` parameter in the `preprocess` method. + image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`): + Mean to use if normalizing the image. This is a float or list of floats the length of the number of + channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method. Can be + overridden by the `image_mean` parameter in the `preprocess` method. + image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`): + Standard deviation to use if normalizing the image. This is a float or list of floats the length of the + number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method. + Can be overridden by the `image_std` parameter in the `preprocess` method. + do_convert_rgb (`bool`, *optional*, defaults to `True`): + Whether to convert the image to RGB. + """ + + model_input_names = ["pixel_values"] + + def __init__( + self, + do_resize: bool = True, + size: Dict[str, int] = None, + resample: PILImageResampling = PILImageResampling.BICUBIC, + do_rescale: bool = True, + rescale_factor: Union[int, float] = 1 / 255, + do_normalize: bool = True, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + do_convert_rgb: bool = True, + **kwargs, + ) -> None: + super().__init__(**kwargs) + size = size if size is not None else {"height": 384, "width": 384} + size = get_size_dict(size, default_to_square=True) + + self.do_resize = do_resize + self.size = size + self.resample = resample + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.do_normalize = do_normalize + self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN + self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD + self.do_convert_rgb = do_convert_rgb + + def resize( + self, + image: np.ndarray, + size: Dict[str, int], + resample: PILImageResampling = PILImageResampling.BICUBIC, + data_format: Optional[Union[str, ChannelDimension]] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, + **kwargs, + ) -> np.ndarray: + """ + Resize an image to `(size["height"], size["width"])`. + + Args: + image (`np.ndarray`): + Image to resize. + size (`Dict[str, int]`): + Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image. + resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`): + `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BICUBIC`. + data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the output image. If unset, the channel dimension format of the input + image is used. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the input image. If unset, the channel dimension format is inferred + from the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. + + Returns: + `np.ndarray`: The resized image. + """ + size = get_size_dict(size) + if "height" not in size or "width" not in size: + raise ValueError(f"The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}") + output_size = (size["height"], size["width"]) + return resize( + image, + size=output_size, + resample=resample, + data_format=data_format, + input_data_format=input_data_format, + **kwargs, + ) + + @filter_out_non_signature_kwargs() + def preprocess( + self, + images: ImageInput, + do_resize: Optional[bool] = None, + size: Optional[Dict[str, int]] = None, + resample: PILImageResampling = None, + do_rescale: Optional[bool] = None, + rescale_factor: Optional[float] = None, + do_normalize: Optional[bool] = None, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + do_convert_rgb: bool = None, + data_format: ChannelDimension = ChannelDimension.FIRST, + input_data_format: Optional[Union[str, ChannelDimension]] = None, + ) -> PIL.Image.Image: + """ + Preprocess an image or batch of images. + + Args: + images (`ImageInput`): + Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If + passing in images with pixel values between 0 and 1, set `do_rescale=False`. + do_resize (`bool`, *optional*, defaults to `self.do_resize`): + Whether to resize the image. + size (`Dict[str, int]`, *optional*, defaults to `self.size`): + Controls the size of the image after `resize`. The shortest edge of the image is resized to + `size["shortest_edge"]` whilst preserving the aspect ratio. If the longest edge of this resized image + is > `int(size["shortest_edge"] * (1333 / 800))`, then the image is resized again to make the longest + edge equal to `int(size["shortest_edge"] * (1333 / 800))`. + resample (`PILImageResampling`, *optional*, defaults to `self.resample`): + Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`. + do_rescale (`bool`, *optional*, defaults to `self.do_rescale`): + Whether to rescale the image values between [0 - 1]. + rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`): + Rescale factor to rescale the image by if `do_rescale` is set to `True`. + do_normalize (`bool`, *optional*, defaults to `self.do_normalize`): + Whether to normalize the image. + image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`): + Image mean to normalize the image by if `do_normalize` is set to `True`. + image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`): + Image standard deviation to normalize the image by if `do_normalize` is set to `True`. + do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`): + Whether to convert the image to RGB. + return_tensors (`str` or `TensorType`, *optional*): + The type of tensors to return. Can be one of: + - Unset: Return a list of `np.ndarray`. + - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`. + - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`. + - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`. + - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`. + data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`): + The channel dimension format for the output image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + - Unset: Use the channel dimension format of the input image. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the input image. If unset, the channel dimension format is inferred + from the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. + """ + do_resize = do_resize if do_resize is not None else self.do_resize + resample = resample if resample is not None else self.resample + do_rescale = do_rescale if do_rescale is not None else self.do_rescale + rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor + do_normalize = do_normalize if do_normalize is not None else self.do_normalize + image_mean = image_mean if image_mean is not None else self.image_mean + image_std = image_std if image_std is not None else self.image_std + do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb + + size = size if size is not None else self.size + size = get_size_dict(size, default_to_square=False) + images = make_flat_list_of_images(images) + + if not valid_images(images): + raise ValueError( + "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, " + "torch.Tensor, tf.Tensor or jax.ndarray." + ) + + validate_preprocess_arguments( + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + do_resize=do_resize, + size=size, + resample=resample, + ) + # PIL RGBA images are converted to RGB + if do_convert_rgb: + images = [convert_to_rgb(image) for image in images] + + # All transformations expect numpy arrays. + images = [to_numpy_array(image) for image in images] + + if do_rescale and is_scaled_image(images[0]): + logger.warning_once( + "It looks like you are trying to rescale already rescaled images. If the input" + " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again." + ) + + if input_data_format is None: + # We assume that all images have the same channel dimension format. + input_data_format = infer_channel_dimension_format(images[0]) + + if do_resize: + images = [ + self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format) + for image in images + ] + + if do_rescale: + images = [ + self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format) + for image in images + ] + + if do_normalize: + images = [ + self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format) + for image in images + ] + + images = [ + to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images + ] + + encoded_outputs = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors) + + return encoded_outputs + + def crop_image_to_patches( + self, + image: ImageInput, + min_patches: int, + max_patches: int, + use_thumbnail: bool = True, + patch_size: Union[Tuple, int, dict] = None, + return_numpy: bool = False, + data_format: ChannelDimension = None, + ): + """ + Crop the image to patches and return a list of cropped images. + The number of patches and their grid arrangement are determined by the original image size, + the target patch size and the minimum and maximum number of patches. + The aspect ratio of the patches grid is chosen to be the closest to the original image aspect ratio. + + Args: + image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`): + The image to be cropped. The image can be a PIL image, NumPy array or PyTorch tensor. + min_patches (`int`): + The minimum number of patches to be extracted from the image. + max_patches (`int`): + The maximum number of patches to be extracted from the image. + use_thumbnail (`bool`, *optional*, defaults to `True`): + Whether to add a thumbnail image to the list of cropped patches. + patch_size (`int`, `Tuple[int, int]`, `dict`, *optional*): + The size of the output patches. + return_numpy (`bool`, *optional*, defaults to `False`): + Whether to return the cropped images as NumPy arrays. + data_format (`ChannelDimension`, *optional*): + The format of the image data. If `None`, the format is inferred from the input image. + + Returns: + List[`PIL.Image.Image`] or List[np.ndarray]: The list of cropped images. + """ + patch_size = patch_size if patch_size is not None else self.size + patch_size = get_size_dict(patch_size, default_to_square=True) + original_size = get_size_dict(image.size, height_width_order=False) + do_rescale = False + if not isinstance(image, PIL.Image.Image): + do_rescale = _rescale_for_pil_conversion(image) + image = to_pil_image(image, do_rescale=do_rescale) + + patch_size_height, patch_size_width = patch_size["height"], patch_size["width"] + original_height, original_width = original_size["height"], original_size["width"] + # find the closest aspect ratio to the target + num_columns, num_rows = get_optimal_tiled_canvas( + (original_height, original_width), (patch_size_height, patch_size_width), min_patches, max_patches + ) + + # calculate the target width and height + target_width = patch_size_width * num_columns + target_height = patch_size_height * num_rows + num_blocks = num_columns * num_rows + + # resize the image so that each patch is of patch_size + resized_image = image.resize((target_width, target_height)) + + # split the image into patches + processed_images = [] + for i in range(num_blocks): + column = i % num_columns + row = i // num_columns + box = ( + column * patch_size_width, + row * patch_size_height, + (column + 1) * patch_size_width, + (row + 1) * patch_size_height, + ) + # split the image + patch_image = resized_image.crop(box) + processed_images.append(patch_image) + + if use_thumbnail and len(processed_images) != 1: + thumbnail_img = image.resize((patch_size_width, patch_size_height)) + processed_images.append(thumbnail_img) + + if return_numpy: + processed_images_numpy = [] + for processed_image in processed_images: + processed_image = np.array(processed_image) + # If the input image channel dimension was of size 1, then it is dropped when converting to a PIL image + # so we need to add it back if necessary. + processed_image = ( + np.expand_dims(processed_image, axis=-1) if processed_image.ndim == 2 else processed_image + ) + # The image is always in channels last format after converting from a PIL image + if data_format is not None: + processed_image = to_channel_dimension_format( + processed_image, data_format, input_channel_dim=ChannelDimension.LAST + ) + # If an image was rescaled to be in the range [0, 255] before converting to a PIL image, then we need to + # rescale it back to the original range. + processed_image = self.rescale(processed_image, 1 / 255) if do_rescale else processed_image + processed_images_numpy.append(processed_image) + processed_images = processed_images_numpy + + return processed_images + + +__all__ = ["GotOcr2ImageProcessor"] diff --git a/src/transformers/models/got_ocr2/modeling_got_ocr2.py b/src/transformers/models/got_ocr2/modeling_got_ocr2.py new file mode 100644 index 0000000000..957e05bea7 --- /dev/null +++ b/src/transformers/models/got_ocr2/modeling_got_ocr2.py @@ -0,0 +1,1005 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/got_ocr2/modular_got_ocr2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_got_ocr2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2024 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. + + +import collections +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...activations import ACT2FN +from ...generation import GenerationMixin +from ...modeling_outputs import ModelOutput +from ...modeling_utils import PreTrainedModel +from ...utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + replace_return_docstrings, +) +from ..auto import AutoModelForCausalLM +from .configuration_got_ocr2 import GotOcr2Config, GotOcr2VisionConfig + + +_CONFIG_FOR_DOC = "GotOcr2Config" + + +class GotOcr2MLPBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.lin1 = nn.Linear(config.hidden_size, config.mlp_dim) + self.lin2 = nn.Linear(config.mlp_dim, config.hidden_size) + self.act = ACT2FN[config.hidden_act] + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.lin1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.lin2(hidden_states) + return hidden_states + + +class GotOcr2VisionAttention(nn.Module): + """Multi-head Attention block with relative position embeddings.""" + + def __init__(self, config, window_size): + super().__init__() + input_size = ( + (config.image_size // config.patch_size, config.image_size // config.patch_size) + if window_size == 0 + else (window_size, window_size) + ) + + self.num_attention_heads = config.num_attention_heads + head_dim = config.hidden_size // config.num_attention_heads + self.scale = head_dim**-0.5 + self.dropout = config.attention_dropout + + self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.qkv_bias) + self.proj = nn.Linear(config.hidden_size, config.hidden_size) + + self.use_rel_pos = config.use_rel_pos + if self.use_rel_pos: + if input_size is None: + raise ValueError("Input size must be provided if using relative positional encoding.") + + # initialize relative positional embeddings + self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) + self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim)) + + def get_rel_pos(self, q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor: + """ + Get relative positional embeddings according to the relative positions of + query and key sizes. + + Args: + q_size (int): + size of the query. + k_size (int): + size of key k. + rel_pos (`torch.Tensor`): + relative position embeddings (L, channel). + + Returns: + Extracted positional embeddings according to relative positions. + """ + max_rel_dist = int(2 * max(q_size, k_size) - 1) + # Interpolate rel pos. + rel_pos_resized = F.interpolate( + rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), + size=max_rel_dist, + mode="linear", + ) + rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0) + + # Scale the coords with short length if shapes for q and k are different. + q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0) + k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0) + relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) + + return rel_pos_resized[relative_coords.long()] + + def add_decomposed_rel_pos( + self, + attn: torch.Tensor, + query: torch.Tensor, + rel_pos_h: torch.Tensor, + rel_pos_w: torch.Tensor, + q_size: Tuple[int, int], + k_size: Tuple[int, int], + ) -> torch.Tensor: + """ + Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. + https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py + + Args: + attn (`torch.Tensor`): + attention map. + query (`torch.Tensor`): + query q in the attention layer with shape (batch_size, query_height * query_width, channel). + rel_pos_h (`torch.Tensor`): + relative position embeddings (Lh, channel) for height axis. + rel_pos_w (`torch.Tensor`): + relative position embeddings (Lw, channel) for width axis. + q_size (tuple): + spatial sequence size of query q with (query_height, query_width). + k_size (tuple): + spatial sequence size of key k with (key_height, key_width). + + Returns: + attn (`torch.Tensor`): + attention map with added relative positional embeddings. + """ + query_height, query_width = q_size + key_height, key_width = k_size + relative_position_height = self.get_rel_pos(query_height, key_height, rel_pos_h) + relative_position_width = self.get_rel_pos(query_width, key_width, rel_pos_w) + + batch_size, _, dim = query.shape + reshaped_query = query.reshape(batch_size, query_height, query_width, dim) + rel_h = torch.einsum("bhwc,hkc->bhwk", reshaped_query, relative_position_height) + rel_w = torch.einsum("bhwc,wkc->bhwk", reshaped_query, relative_position_width) + attn = attn.reshape(batch_size, query_height, query_width, key_height, key_width) + attn = attn + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :] + attn = attn.reshape(batch_size, query_height * query_width, key_height * key_width) + return attn + + def forward(self, hidden_states: torch.Tensor, output_attentions=False) -> torch.Tensor: + batch_size, height, width, _ = hidden_states.shape + # qkv with shape (3, batch_size, nHead, height * width, channel) + qkv = ( + self.qkv(hidden_states) + .reshape(batch_size, height * width, 3, self.num_attention_heads, -1) + .permute(2, 0, 3, 1, 4) + ) + # q, k, v with shape (batch_size * nHead, height * width, channel) + query, key, value = qkv.reshape(3, batch_size * self.num_attention_heads, height * width, -1).unbind(0) + + attn_weights = (query * self.scale) @ key.transpose(-2, -1) + + if self.use_rel_pos: + attn_weights = self.add_decomposed_rel_pos( + attn_weights, query, self.rel_pos_h, self.rel_pos_w, (height, width), (height, width) + ) + + attn_weights = torch.nn.functional.softmax(attn_weights, dtype=torch.float32, dim=-1).to(query.dtype) + + attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + + attn_output = (attn_probs @ value).reshape(batch_size, self.num_attention_heads, height, width, -1) + attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(batch_size, height, width, -1) + + attn_output = self.proj(attn_output) + + if output_attentions: + outputs = (attn_output, attn_weights) + else: + outputs = (attn_output, None) + + return outputs + + +class GotOcr2VisionLayer(nn.Module): + def __init__(self, config, window_size): + super().__init__() + self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.attn = GotOcr2VisionAttention(config, window_size) + self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.mlp = GotOcr2MLPBlock(config) + self.window_size = window_size + + def window_partition(self, hidden_states: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]: + """ + Args: + Partition into non-overlapping windows with padding if needed. + hidden_states (tensor): input tokens with [batch_size, height, width, channel]. window_size (int): window + size. + + Returns: + windows: windows after partition with [batch_size * num_windows, window_size, window_size, channel]. + (pad_height, pad_width): padded height and width before partition + """ + batch_size, height, width, channel = hidden_states.shape + + pad_h = (window_size - height % window_size) % window_size + pad_w = (window_size - width % window_size) % window_size + hidden_states = F.pad(hidden_states, (0, 0, 0, pad_w, 0, pad_h)) + pad_height, pad_width = height + pad_h, width + pad_w + + hidden_states = hidden_states.reshape( + batch_size, pad_height // window_size, window_size, pad_width // window_size, window_size, channel + ) + windows = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(-1, window_size, window_size, channel) + return windows, (pad_height, pad_width) + + def window_unpartition( + self, windows: torch.Tensor, window_size: int, padding_shape: Tuple[int, int], original_shape: Tuple[int, int] + ) -> torch.Tensor: + """ + Args: + Window unpartition into original sequences and removing padding. + hidden_states (tensor): + input tokens with [batch_size * num_windows, window_size, window_size, channel]. + window_size (int): + window size. + padding_shape (Tuple): + padded height and width (pad_height, pad_width). + original_shape (Tuple): original height and width (height, width) before padding. + + Returns: + hidden_states: unpartitioned sequences with [batch_size, height, width, channel]. + """ + pad_height, pad_width = padding_shape + height, width = original_shape + batch_size = windows.shape[0] // (pad_height * pad_width // window_size // window_size) + hidden_states = windows.reshape( + batch_size, pad_height // window_size, pad_width // window_size, window_size, window_size, -1 + ) + hidden_states = ( + hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(batch_size, pad_height, pad_width, -1) + ) + + hidden_states = hidden_states[:, :height, :width, :].contiguous() + return hidden_states + + def forward( + self, + hidden_states: torch.Tensor, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.FloatTensor]: + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + # Window partition + if self.window_size > 0: + height, width = hidden_states.shape[1], hidden_states.shape[2] + hidden_states, padding_shape = self.window_partition(hidden_states, self.window_size) + + hidden_states, attn_weights = self.attn( + hidden_states=hidden_states, + output_attentions=output_attentions, + ) + # Reverse window partition + if self.window_size > 0: + hidden_states = self.window_unpartition(hidden_states, self.window_size, padding_shape, (height, width)) + + hidden_states = residual + hidden_states + layernorm_output = self.layer_norm2(hidden_states) + hidden_states = hidden_states + self.mlp(layernorm_output) + + outputs = (hidden_states,) + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +@dataclass +class GotOcr2VisionEncoderOutput(ModelOutput): + """ + Base class for got_ocr2 vision model's outputs that also contains image embeddings obtained by applying the projection + layer to the pooler_output. + + Args: + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The image embeddings obtained by applying the projection layer to the pooler_output. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + image_embeds: Optional[torch.FloatTensor] = None + last_hidden_state: torch.FloatTensor = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +class GotOcr2PatchEmbeddings(nn.Module): + """ + This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial + `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a + Transformer. + """ + + def __init__(self, config): + super().__init__() + image_size, patch_size = config.image_size, config.patch_size + num_channels, hidden_size = config.num_channels, config.hidden_size + image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) + patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) + num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + + self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) + + def forward(self, pixel_values): + batch_size, num_channels, height, width = pixel_values.shape + if num_channels != self.num_channels: + raise ValueError( + "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + ) + if height != self.image_size[0] or width != self.image_size[1]: + raise ValueError( + f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})." + ) + embeddings = self.projection(pixel_values).permute(0, 2, 3, 1) + return embeddings + + +class GotOcr2LayerNorm(nn.Module): + r"""LayerNorm that supports two data formats: channels_last (default) or channels_first. + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, + width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width). + """ + + def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + if self.data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError(f"Unsupported data format: {self.data_format}") + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + x = torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + elif self.data_format == "channels_first": + input_dtype = x.dtype + x = x.float() + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = x.to(dtype=input_dtype) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x + + +class GotOcr2VisionNeck(nn.Module): + def __init__(self, config: GotOcr2VisionConfig): + super().__init__() + self.config = config + + self.conv1 = nn.Conv2d(config.hidden_size, config.output_channels, kernel_size=1, bias=False) + self.layer_norm1 = GotOcr2LayerNorm(config.output_channels, data_format="channels_first") + self.conv2 = nn.Conv2d(config.output_channels, config.output_channels, kernel_size=3, padding=1, bias=False) + self.layer_norm2 = GotOcr2LayerNorm(config.output_channels, data_format="channels_first") + + def forward(self, hidden_states): + hidden_states = hidden_states.permute(0, 3, 1, 2) + hidden_states = self.conv1(hidden_states) + hidden_states = self.layer_norm1(hidden_states) + + hidden_states = self.conv2(hidden_states) + hidden_states = self.layer_norm2(hidden_states) + return hidden_states + + +class GotOcr2VisionEncoder(nn.Module): + def __init__(self, config: GotOcr2VisionConfig): + super().__init__() + self.config = config + self.image_size = config.image_size + + self.patch_embed = GotOcr2PatchEmbeddings(config) + + self.pos_embed = None + if config.use_abs_pos: + # Initialize absolute positional embedding with pretrain image size. + self.pos_embed = nn.Parameter( + torch.zeros( + 1, + config.image_size // config.patch_size, + config.image_size // config.patch_size, + config.hidden_size, + ) + ) + + self.layers = nn.ModuleList() + for i in range(config.num_hidden_layers): + layer = GotOcr2VisionLayer( + config, + window_size=config.window_size if i not in config.global_attn_indexes else 0, + ) + self.layers.append(layer) + + self.neck = GotOcr2VisionNeck(config) + + self.gradient_checkpointing = False + + def get_input_embeddings(self): + return self.patch_embed + + def forward( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, GotOcr2VisionEncoderOutput]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + hidden_states = self.patch_embed(pixel_values) + if self.pos_embed is not None: + hidden_states = hidden_states + self.pos_embed + + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + + for i, layer_module in enumerate(self.layers): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + layer_module.__call__, + hidden_states, + ) + else: + layer_outputs = layer_module(hidden_states, output_attentions=output_attentions) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + hidden_states = self.neck(hidden_states) + + if not return_dict: + outputs = (hidden_states,) + if output_hidden_states: + outputs = outputs + (all_hidden_states,) + if output_attentions: + outputs = outputs + (all_self_attentions,) + return outputs + + return GotOcr2VisionEncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +class GotOcr2MultiModalProjector(nn.Module): + def __init__(self, config: GotOcr2Config): + super().__init__() + vision_output_channels = config.vision_config.output_channels + language_hidden_size = config.text_config.hidden_size + self.conv_upsampler1 = nn.Conv2d( + vision_output_channels, vision_output_channels * 2, kernel_size=3, stride=2, padding=1, bias=False + ) + self.conv_upsampler2 = nn.Conv2d( + vision_output_channels * 2, language_hidden_size, kernel_size=3, stride=2, padding=1, bias=False + ) + self.multimodal_projector = nn.Linear(language_hidden_size, language_hidden_size) + + def forward(self, vision_embeddings: torch.Tensor) -> torch.Tensor: + hidden_state = self.conv_upsampler1(vision_embeddings) + hidden_state = self.conv_upsampler2(hidden_state) + hidden_state = hidden_state.flatten(2).permute(0, 2, 1) + hidden_state = self.multimodal_projector(hidden_state) + return hidden_state + + +@dataclass +class GotOcr2CausalLMOutputWithPast(ModelOutput): + """ + Base class for GotOcr2 causal language model (or autoregressive) outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size (batch_size, num_images, sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. + """ + + loss: Optional[torch.FloatTensor] = None + logits: torch.FloatTensor = None + past_key_values: Optional[List[torch.FloatTensor]] = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + image_hidden_states: Optional[torch.FloatTensor] = None + + +GOT_OCR2_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`GotOcr2Config`] or [`GotOcr2VisionConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare LLaMA Model outputting raw hidden-states without any specific head on top.", + GOT_OCR2_START_DOCSTRING, +) +class GotOcr2PreTrainedModel(PreTrainedModel): + config_class = GotOcr2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["GotOcr2VisionAttention"] + _skip_keys_device_placement = "past_key_values" + _supports_cache_class = True + _supports_flash_attn_2 = True + _supports_sdpa = True + + def _init_weights(self, module): + # important: this ported version of GotOcr2 isn't meant for training from scratch - only + # inference and fine-tuning - so the proper init weights code has been removed - the original codebase + # https://github.com/haotian-liu/GotOcr2/tree/main/got_ocr2 should serve for that purpose + std = ( + self.config.initializer_range + if hasattr(self.config, "initializer_range") + else self.config.text_config.initializer_range + ) + + if hasattr(module, "class_embedding"): + module.class_embedding.data.normal_(mean=0.0, std=std) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +GOT_OCR2_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)): + The tensors corresponding to the input images. Pixel values can be obtained using + [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details ([]`GotOcr2Processor`] uses + [`CLIPImageProcessor`] for processing images). + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids) + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + vision_feature_layer (`Union[int, List[int]], *optional*, defaults to -2`): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*, defaults to `"default"`): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"`. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + + +@add_start_docstrings( + """The GOT_OCR2 model which consists of a vision backbone and a language model.""", + GOT_OCR2_START_DOCSTRING, +) +class GotOcr2ForConditionalGeneration(GotOcr2PreTrainedModel, GenerationMixin): + def __init__(self, config: GotOcr2Config): + super().__init__(config) + self.vision_tower = GotOcr2VisionEncoder(config.vision_config) + + self.multi_modal_projector = GotOcr2MultiModalProjector(config) + self.vocab_size = config.text_config.vocab_size + self.language_model = AutoModelForCausalLM.from_config(config.text_config) + + if self.language_model._tied_weights_keys is not None: + self._tied_weights_keys = [f"language_model.{k}" for k in self.language_model._tied_weights_keys] + + self.pad_token_id = config.pad_token_id + + self.post_init() + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.language_model.get_output_embeddings() + + def set_output_embeddings(self, new_embeddings): + self.language_model.set_output_embeddings(new_embeddings) + + def set_decoder(self, decoder): + self.language_model.set_decoder(decoder) + + def get_decoder(self): + return self.language_model.get_decoder() + + def get_image_features( + self, + pixel_values: torch.FloatTensor, + ): + """ + Obtains image last hidden states from the vision tower and apply multimodal projection. + + Args: + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) + Returns: + image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + """ + image_outputs = self.vision_tower(pixel_values).last_hidden_state + return self.multi_modal_projector(image_outputs) + + def _merge_input_ids_with_image_features(self, image_features, inputs_embeds, input_ids, attention_mask, labels): + num_images, num_image_patches, embed_dim = image_features.shape + batch_size, sequence_length = input_ids.shape + left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(self.pad_token_id)) + # 1. Create a mask to know where special image tokens are + special_image_token_mask = input_ids == self.config.image_token_index + num_special_image_tokens = torch.sum(special_image_token_mask, dim=-1) + # Compute the maximum embed dimension + max_embed_dim = (num_special_image_tokens.max() * (num_image_patches - 1)) + sequence_length + batch_indices, non_image_indices = torch.where(input_ids != self.config.image_token_index) + + # 2. Compute the positions where text should be written + # Calculate new positions for text tokens in merged image-text sequence. + # `special_image_token_mask` identifies image tokens. Each image token will be replaced by `nb_text_tokens_per_images - 1` text tokens. + # `torch.cumsum` computes how each image token shifts subsequent text token positions. + # - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one. + new_token_positions = torch.cumsum((special_image_token_mask * (num_image_patches - 1) + 1), -1) - 1 + nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] + if left_padding: + new_token_positions += nb_image_pad[:, None] # offset for left padding + text_to_overwrite = new_token_positions[batch_indices, non_image_indices] + + # 3. Create the full embedding, already padded to the maximum position + final_embedding = torch.zeros( + batch_size, max_embed_dim, embed_dim, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + final_attention_mask = torch.zeros( + batch_size, max_embed_dim, dtype=attention_mask.dtype, device=inputs_embeds.device + ) + if labels is not None: + final_labels = torch.full( + (batch_size, max_embed_dim), self.config.ignore_index, dtype=input_ids.dtype, device=input_ids.device + ) + # In case the Vision model or the Language model has been offloaded to CPU, we need to manually + # set the corresponding tensors into their correct target device. + target_device = inputs_embeds.device + batch_indices, non_image_indices, text_to_overwrite = ( + batch_indices.to(target_device), + non_image_indices.to(target_device), + text_to_overwrite.to(target_device), + ) + attention_mask = attention_mask.to(target_device) + + # 4. Fill the embeddings based on the mask. If we have ["hey" "", "how", "are"] + # we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the image features + final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices] + final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices] + if labels is not None: + final_labels[batch_indices, text_to_overwrite] = labels[batch_indices, non_image_indices] + + # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) + image_to_overwrite = torch.full( + (batch_size, max_embed_dim), True, dtype=torch.bool, device=inputs_embeds.device + ) + image_to_overwrite[batch_indices, text_to_overwrite] = False + if left_padding: + image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device) + else: + mask = torch.ones_like(image_to_overwrite, dtype=torch.bool).cumsum(-1) - 1 + padding_mask = mask <= new_token_positions[:, -1:].to(target_device) + image_to_overwrite &= padding_mask + + if image_to_overwrite.sum() != image_features.shape[:-1].numel(): + raise ValueError( + f"The input provided to the model are wrong. The number of image tokens is {torch.sum(special_image_token_mask)} while" + f" the number of image given to the model is {num_images}. This prevents correct indexing and breaks batch generation." + ) + + final_embedding[image_to_overwrite] = image_features.contiguous().reshape(-1, embed_dim).to(target_device) + final_attention_mask |= image_to_overwrite + position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1) + + # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. + batch_indices, pad_indices = torch.where(input_ids == self.pad_token_id) + indices_to_mask = new_token_positions[batch_indices, pad_indices] + + final_embedding[batch_indices, indices_to_mask] = 0 + + if labels is None: + final_labels = None + + return final_embedding, final_attention_mask, final_labels, position_ids + + @add_start_docstrings_to_model_forward(GOT_OCR2_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=GotOcr2CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + ) -> Union[Tuple, GotOcr2CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + logits_to_keep (`int` or `torch.Tensor`, *optional*): + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. + This is useful when using packed tensor format (single dimension for batch and sequence length). + + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, GotOcr2ForConditionalGeneration, TextStreamer + + >>> model = GotOcr2ForConditionalGeneration.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf").to("cuda") + >>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + + >>> url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(image, return_tensors="pt", color="green").to("cuda") + + >>> # Generate + >>> streamer = TextStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True) + >>> generate_ids = model.generate( + ... **inputs, + ... do_sample=False, + ... tokenizer = processor.tokenizer, + ... stop_strings='<|im_end|>', + ... streamer=streamer, + ... max_new_tokens=4096, + ... ) + "You should keep in mind what features from the module should be used, especially + when you're planning to sell a template." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if pixel_values is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" + ) + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + image_features = self.get_image_features(pixel_values=pixel_values.to(inputs_embeds.dtype)) + n_image_tokens = (input_ids == self.config.image_token_index).sum() + n_image_features = image_features.shape[0] * image_features.shape[1] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1) + special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device) + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + logits_to_keep=logits_to_keep, + ) + + logits = outputs[0] + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + if attention_mask is not None: + # we use the input attention mask to shift the logits and labels, because it is 2D. + # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft + shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(logits.device) + shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous() + shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() + else: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1).to(shift_logits.device) + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return GotOcr2CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features if pixel_values is not None else None, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + attention_mask=None, + cache_position=None, + logits_to_keep=None, + **kwargs, + ): + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model + + model_inputs = self.language_model.prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + logits_to_keep=logits_to_keep, + **kwargs, + ) + + if cache_position[0] == 0: + # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore + # Otherwise we need pixel values to be passed to model + model_inputs["pixel_values"] = pixel_values + + return model_inputs + + +__all__ = ["GotOcr2PreTrainedModel", "GotOcr2ForConditionalGeneration"] diff --git a/src/transformers/models/got_ocr2/modular_got_ocr2.py b/src/transformers/models/got_ocr2/modular_got_ocr2.py new file mode 100644 index 0000000000..899075683e --- /dev/null +++ b/src/transformers/models/got_ocr2/modular_got_ocr2.py @@ -0,0 +1,984 @@ +# coding=utf-8 +# Copyright 2024 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. + + +from functools import lru_cache +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +import torch.utils.checkpoint + +from transformers.models.blip.image_processing_blip import BlipImageProcessor +from transformers.models.llava.modeling_llava import ( + LlavaCausalLMOutputWithPast, + LlavaForConditionalGeneration, + LlavaPreTrainedModel, +) +from transformers.models.sam.modeling_sam import SamMLPBlock, SamVisionAttention, SamVisionEncoder, SamVisionLayer +from transformers.processing_utils import ImagesKwargs, ProcessingKwargs, ProcessorMixin, TextKwargs, Unpack +from transformers.tokenization_utils_base import ( + PreTokenizedInput, + TextInput, +) + +from ...configuration_utils import PretrainedConfig +from ...image_processing_utils import BatchFeature, get_size_dict +from ...image_transforms import ( + _rescale_for_pil_conversion, + to_channel_dimension_format, + to_pil_image, +) +from ...image_utils import ChannelDimension, ImageInput +from ...utils import ( + add_start_docstrings_to_model_forward, + is_vision_available, + logging, + replace_return_docstrings, +) +from ..auto import CONFIG_MAPPING, AutoConfig, AutoModelForCausalLM + + +if is_vision_available(): + import PIL + + from ...image_utils import load_images + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "GotOcr2Config" + + +class GotOcr2VisionConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`GotOcr2VisionModel`]. It is used to instantiate a GOT_OCR2 + vision encoder according to the specified arguments, defining the model architecture. Instantiating a configuration + defaults will yield a similar configuration to that of the SAM ViT-h + [facebook/sam-vit-huge](https://huggingface.co/facebook/sam-vit-huge) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + hidden_size (`int`, *optional*, defaults to 768): + Dimensionality of the encoder layers and the pooler layer. + output_channels (`int`, *optional*, defaults to 256): + Dimensionality of the output channels in the Patch Encoder. + num_hidden_layers (`int`, *optional*, defaults to 12): + Number of hidden layers in the Transformer encoder. + num_attention_heads (`int`, *optional*, defaults to 12): + Number of attention heads for each attention layer in the Transformer encoder. + num_channels (`int`, *optional*, defaults to 3): + Number of channels in the input image. + image_size (`int`, *optional*, defaults to 1024): + Expected resolution. Target size of the resized input image. + patch_size (`int`, *optional*, defaults to 16): + Size of the patches to be extracted from the input image. + hidden_act (`str`, *optional*, defaults to `"gelu"`): + The non-linear activation function (function or string) + layer_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the layer normalization layers. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + initializer_range (`float`, *optional*, defaults to 1e-10): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + qkv_bias (`bool`, *optional*, defaults to `True`): + Whether to add a bias to query, key, value projections. + use_abs_pos (`bool`, *optional*, defaults to `True`): + Whether to use absolute position embedding. + use_rel_pos (`bool`, *optional*, defaults to `True`): + Whether to use relative position embedding. + window_size (`int`, *optional*, defaults to 14): + Window size for relative position. + global_attn_indexes (`List[int]`, *optional*, defaults to `[2, 5, 8, 11]`): + The indexes of the global attention layers. + mlp_dim (`int`, *optional*, defaults to 3072): + The dimensionality of the MLP layer in the Transformer encoder. + """ + + base_config_key = "vision_config" + + def __init__( + self, + hidden_size=768, + output_channels=256, + num_hidden_layers=12, + num_attention_heads=12, + num_channels=3, + image_size=1024, + patch_size=16, + hidden_act="gelu", + layer_norm_eps=1e-06, + attention_dropout=0.0, + initializer_range=1e-10, + qkv_bias=True, + use_abs_pos=True, + use_rel_pos=True, + window_size=14, + global_attn_indexes=[2, 5, 8, 11], + mlp_dim=3072, + **kwargs, + ): + super().__init__(**kwargs) + + self.hidden_size = hidden_size + self.output_channels = output_channels + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.hidden_act = hidden_act + self.layer_norm_eps = layer_norm_eps + self.attention_dropout = attention_dropout + self.initializer_range = initializer_range + self.qkv_bias = qkv_bias + self.use_abs_pos = use_abs_pos + self.use_rel_pos = use_rel_pos + self.window_size = window_size + self.global_attn_indexes = global_attn_indexes + self.mlp_dim = mlp_dim + + +class GotOcr2Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`GotOcr2ForConditionalGeneration`]. It is used to instantiate a + GotOcr2 model according to the specified arguments, defining the model architecture. Instantiating a configuration + with the defaults will yield a similar configuration to that of GOT-OCR-2.0. + + e.g [stepfun-ai/GOT-OCR-2.0-hf](https://huggingface.co/stepfun-ai/GOT-OCR-2.0-hf) + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vision_config (`Union[AutoConfig, dict]`, *optional*, defaults to `CLIPVisionConfig`): + The config object or dictionary of the vision backbone. + text_config (`Union[AutoConfig, dict]`, *optional*, defaults to `LlamaConfig`): + The config object or dictionary of the text backbone. + ignore_index (`int`, *optional*, defaults to -100): + The ignore index for the loss function. + image_token_index (`int`, *optional*, defaults to 151859): + The image token index to encode the image prompt. + image_seq_length (`int`, *optional*, defaults to 576): + Sequence length of one image embedding. + pad_token_id (`int`, *optional*, defaults to -1): + Padding token id. + + ```python + >>> from transformers import GotOcr2ForConditionalGeneration, GotOcr2Config + + >>> # Initializing a GotOcr2 style configuration + >>> configuration = GotOcr2Config() + + >>> # Initializing a model from the Qwen2-VL-7B style configuration + >>> model = GotOcr2ForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "got_ocr2" + sub_configs = {"text_config": AutoConfig, "vision_config": GotOcr2VisionConfig} + + def __init__( + self, + vision_config=None, + text_config=None, + ignore_index=-100, + image_token_index=151859, + image_seq_length=576, + pad_token_id=-1, + **kwargs, + ): + self.ignore_index = ignore_index + self.image_token_index = image_token_index + self.image_seq_length = image_seq_length + self.pad_token_id = pad_token_id + + if vision_config is None: + self.vision_config = GotOcr2VisionConfig() + elif isinstance(vision_config, dict): + self.vision_config = GotOcr2VisionConfig(**vision_config) + elif isinstance(vision_config, GotOcr2VisionConfig): + self.vision_config = vision_config + + if isinstance(text_config, dict): + text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "qwen2" + text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config) + elif text_config is None: + text_config = CONFIG_MAPPING["qwen2"]( + vocab_size=151860, + hidden_size=1024, + intermediate_size=2816, + num_hidden_layers=24, + num_attention_heads=16, + num_key_value_heads=16, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=True, + rope_theta=1000000.0, + rope_scaling=None, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=21, + attention_dropout=0.0, + ) + + self.text_config = text_config + + super().__init__(**kwargs) + + +__all__ = ["GotOcr2VisionConfig", "GotOcr2Config"] + + +class GotOcr2TextKwargs(TextKwargs, total=False): + format: Optional[bool] + + +class GotOcr2ImagesKwargs(ImagesKwargs, total=False): + box: Optional[Union[List, Tuple[float, float], Tuple[float, float, float, float]]] + color: Optional[str] + num_image_tokens: Optional[int] + multi_page: Optional[bool] + crop_to_patches: Optional[bool] + min_patches: Optional[int] + max_patches: Optional[int] + + +class GotOcr2ProcessorKwargs(ProcessingKwargs, total=False): + text_kwargs: GotOcr2TextKwargs + images_kwargs: GotOcr2ImagesKwargs + _defaults = { + "text_kwargs": { + "padding": False, + "format": False, + }, + "images_kwargs": { + "num_image_tokens": 256, + "multi_page": False, + "crop_to_patches": False, + "min_patches": 1, + "max_patches": 12, + }, + } + + +def preprocess_box_annotation(box: Union[List, Tuple], image_size: Tuple[int, int]) -> List: + """ + Convert box annotation to the format [x1, y1, x2, y2] in the range [0, 1000]. + """ + width, height = image_size + if len(box) == 4: + box[0] = int(box[0] / width * 1000) + box[1] = int(box[1] / height * 1000) + box[2] = int(box[2] / width * 1000) + box[3] = int(box[3] / height * 1000) + else: + raise ValueError("Box must be a list or tuple of lists in the form [x1, y1, x2, y2].") + + return list(box) + + +# Similar to image_processing_mllama.get_all_supported_aspect_ratios +@lru_cache(maxsize=10) +def get_all_supported_aspect_ratios(min_image_tiles: int, max_image_tiles: int) -> List[Tuple[int, int]]: + """ + Computes all allowed aspect ratios for a given minimum and maximum number of input tiles. + + This function calculates all possible arrangements of tiles that can be formed + within the constraint of the minimum and maximum number of tiles. Each arrangement is + represented by its aspect ratio (width/height) and the corresponding tile configuration. + + Args: + min_image_tiles (`int`): + The minimum number of tiles allowed. + max_image_tiles (`int`): + The maximum number of tiles allowed. + + Returns: + `List[Tuple[int, int]]`: A list of tuples, each tuple representing a valid (width, height) + configuration in terms of number of tiles. + + Example: + >>> get_all_supported_aspect_ratios(1, 4) + [(1, 1), (1, 2), (2, 1), (1, 3), (3, 1), (1, 4), (2, 2), (4, 1)] + + """ + aspect_ratios = [] + for width in range(1, max_image_tiles + 1): + for height in range(1, max_image_tiles + 1): + if width * height <= max_image_tiles and width * height >= min_image_tiles: + aspect_ratios.append((width, height)) + + aspect_ratios = sorted(aspect_ratios, key=lambda x: x[0] * x[1]) + + return aspect_ratios + + +@lru_cache(maxsize=100) +def get_optimal_tiled_canvas( + original_image_size: Tuple[int, int], + target_tile_size: Tuple[int, int], + min_image_tiles: int, + max_image_tiles: int, +) -> Tuple[int, int]: + """ + Given a minimum and maximum number of tiles, find the canvas with the closest aspect ratio to the + original image aspect ratio. + In case of tie-breaking condition when two canvases have the same aspect ratio difference, we favor the canvas with + more tiles, until the area covered by the tiles is more than twice the target area, in order to avoid unnecessarily + excessive tiling. + """ + possible_tile_arrangements = get_all_supported_aspect_ratios(min_image_tiles, max_image_tiles) + + original_height, original_width = original_image_size + target_tile_height, target_tile_width = target_tile_size + aspect_ratio = original_width / original_height + area = original_width * original_height + + # find the grid with the best aspect ratio + best_ratio_diff = float("inf") + best_grid = (1, 1) + for grid in possible_tile_arrangements: + grid_aspect_ratio = grid[0] / grid[1] + ratio_diff = abs(aspect_ratio - grid_aspect_ratio) + if ratio_diff < best_ratio_diff: + best_ratio_diff = ratio_diff + best_grid = grid + elif ratio_diff == best_ratio_diff: + # if the aspect ratio difference is the same, we favor the grid with more patches + # until the area covered by the patches is more than twice the original image area + if area > 0.5 * target_tile_height * target_tile_width * grid[0] * grid[1]: + best_grid = grid + + return best_grid + + +class GotOcr2ImageProcessor(BlipImageProcessor): + def crop_image_to_patches( + self, + image: ImageInput, + min_patches: int, + max_patches: int, + use_thumbnail: bool = True, + patch_size: Union[Tuple, int, dict] = None, + return_numpy: bool = False, + data_format: ChannelDimension = None, + ): + """ + Crop the image to patches and return a list of cropped images. + The number of patches and their grid arrangement are determined by the original image size, + the target patch size and the minimum and maximum number of patches. + The aspect ratio of the patches grid is chosen to be the closest to the original image aspect ratio. + + Args: + image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`): + The image to be cropped. The image can be a PIL image, NumPy array or PyTorch tensor. + min_patches (`int`): + The minimum number of patches to be extracted from the image. + max_patches (`int`): + The maximum number of patches to be extracted from the image. + use_thumbnail (`bool`, *optional*, defaults to `True`): + Whether to add a thumbnail image to the list of cropped patches. + patch_size (`int`, `Tuple[int, int]`, `dict`, *optional*): + The size of the output patches. + return_numpy (`bool`, *optional*, defaults to `False`): + Whether to return the cropped images as NumPy arrays. + data_format (`ChannelDimension`, *optional*): + The format of the image data. If `None`, the format is inferred from the input image. + + Returns: + List[`PIL.Image.Image`] or List[np.ndarray]: The list of cropped images. + """ + patch_size = patch_size if patch_size is not None else self.size + patch_size = get_size_dict(patch_size, default_to_square=True) + original_size = get_size_dict(image.size, height_width_order=False) + do_rescale = False + if not isinstance(image, PIL.Image.Image): + do_rescale = _rescale_for_pil_conversion(image) + image = to_pil_image(image, do_rescale=do_rescale) + + patch_size_height, patch_size_width = patch_size["height"], patch_size["width"] + original_height, original_width = original_size["height"], original_size["width"] + # find the closest aspect ratio to the target + num_columns, num_rows = get_optimal_tiled_canvas( + (original_height, original_width), (patch_size_height, patch_size_width), min_patches, max_patches + ) + + # calculate the target width and height + target_width = patch_size_width * num_columns + target_height = patch_size_height * num_rows + num_blocks = num_columns * num_rows + + # resize the image so that each patch is of patch_size + resized_image = image.resize((target_width, target_height)) + + # split the image into patches + processed_images = [] + for i in range(num_blocks): + column = i % num_columns + row = i // num_columns + box = ( + column * patch_size_width, + row * patch_size_height, + (column + 1) * patch_size_width, + (row + 1) * patch_size_height, + ) + # split the image + patch_image = resized_image.crop(box) + processed_images.append(patch_image) + + if use_thumbnail and len(processed_images) != 1: + thumbnail_img = image.resize((patch_size_width, patch_size_height)) + processed_images.append(thumbnail_img) + + if return_numpy: + processed_images_numpy = [] + for processed_image in processed_images: + processed_image = np.array(processed_image) + # If the input image channel dimension was of size 1, then it is dropped when converting to a PIL image + # so we need to add it back if necessary. + processed_image = ( + np.expand_dims(processed_image, axis=-1) if processed_image.ndim == 2 else processed_image + ) + # The image is always in channels last format after converting from a PIL image + if data_format is not None: + processed_image = to_channel_dimension_format( + processed_image, data_format, input_channel_dim=ChannelDimension.LAST + ) + # If an image was rescaled to be in the range [0, 255] before converting to a PIL image, then we need to + # rescale it back to the original range. + processed_image = self.rescale(processed_image, 1 / 255) if do_rescale else processed_image + processed_images_numpy.append(processed_image) + processed_images = processed_images_numpy + + return processed_images + + +class GotOcr2Processor(ProcessorMixin): + r""" + Constructs a GotOcr2 processor which wraps a [`GotOcr2ImageProcessor`] and + [`PretrainedTokenizerFast`] tokenizer into a single processor that inherits both the image processor and + tokenizer functionalities. See the [`~GotOcr2Processor.__call__`] and [`~GotOcr2Processor.decode`] for more information. + Args: + image_processor ([`GotOcr2ImageProcessor`], *optional*): + The image processor is a required input. + tokenizer ([`PreTrainedTokenizer`, `PreTrainedTokenizerFast`], *optional*): + The tokenizer is a required input. + chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages + in a chat into a tokenizable string. + """ + + attributes = ["image_processor", "tokenizer"] + valid_kwargs = ["chat_template"] + image_processor_class = "GotOcr2ImageProcessor" + tokenizer_class = "PreTrainedTokenizerFast" + + def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs): + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + self.message_start_token = "<|im_start|>" + self.message_end_token = "<|im_end|>" + self.img_start_token = "" + self.img_end_token = "" + self.img_pad_token = "" + self.system_query = "system\nYou should follow the instructions carefully and explain your answers in detail." + + def _make_list_of_inputs(self, images, text, box, color, multi_page): + if not isinstance(images, (list, tuple)): + images = [images] + if multi_page: + logger.warning("Multi-page inference is enabled but only one image is passed.") + images = [images] + elif isinstance(images[0], (list, tuple)) and not multi_page: + raise ValueError("Nested images are only supported with `multi_page` set to `True`.") + elif not isinstance(images[0], (list, tuple)) and multi_page: + images = [images] + + if isinstance(text, str): + text = [text] + + if not isinstance(box[0], (list, tuple)): + # Use the same box for all images + box = [box for _ in range(len(images))] + if not isinstance(color, (list, tuple)): + color = [color for _ in range(len(images))] + + return images, text, box, color + + def __call__( + self, + images: Optional[ImageInput] = None, + text: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None, + audio=None, + videos=None, + **kwargs: Unpack[GotOcr2ProcessorKwargs], + ) -> BatchFeature: + """ + Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text` + and `kwargs` arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizerFast.__call__`] to encode the text if `text` + is not `None`, otherwise encode default OCR queries which depends on the `format`, `box`, `color`, `multi_page` and + `crop_to_patches` arguments. To prepare the vision inputs, this method forwards the `images` and `kwrags` arguments to + GotOcr2ImageProcessor's [`~GotOcr2ImageProcessor.__call__`] if `images` is not `None`. + + Args: + images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`): + The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch + tensor. Both channels-first and channels-last formats are supported. + text (`str`, `List[str]`, `List[List[str]]`): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + format (`bool`, *optional*): + If set, will add the format token to the query, and the model will return the OCR result with formatting. + box (`List[float]`, `List[Tuple[float, float]]`, `List[Tuple[float, float, float, float]]`, *optional*): + The box annotation to be added to the query. If a list of floats or a tuple of floats is provided, it + will be interpreted as [x1, y1, x2, y2]. If a list of tuples is provided, each tuple should be in the + form (x1, y1, x2, y2). + color (`str`, *optional*): + The color annotation to be added to the query. The model will return the OCR result within the box with + the specified color. + multi_page (`bool`, *optional*): + If set, will enable multi-page inference. The model will return the OCR result across multiple pages. + crop_to_patches (`bool`, *optional*): + If set, will crop the image to patches. The model will return the OCR result upon the patch reference. + min_patches (`int`, *optional*): + The minimum number of patches to be cropped from the image. Only used when `crop_to_patches` is set to + `True`. + max_patches (`int`, *optional*): + The maximum number of patches to be cropped from the image. Only used when `crop_to_patches` is set to + `True`. + + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Acceptable values are: + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + - `'jax'`: Return JAX `jnp.ndarray` objects. + + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + + output_kwargs = self._merge_kwargs( + GotOcr2ProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + format_output = output_kwargs["text_kwargs"].pop("format") + num_image_tokens = output_kwargs["images_kwargs"].pop("num_image_tokens") + box = output_kwargs["images_kwargs"].pop("box", [None]) + color = output_kwargs["images_kwargs"].pop("color", None) + multi_page = output_kwargs["images_kwargs"].pop("multi_page") + crop_to_patches = output_kwargs["images_kwargs"].pop("crop_to_patches") + min_patches = output_kwargs["images_kwargs"].pop("min_patches") + max_patches = output_kwargs["images_kwargs"].pop("max_patches") + + images, text, box, color = self._make_list_of_inputs(images, text, box, color, multi_page) + + # Load images as we need to know the image size + images = load_images(images) + if text is None: + text = [] + for index, (image_group, box_single, color_single) in enumerate(zip(images, box, color)): + if crop_to_patches: + image_group = self.image_processor.crop_image_to_patches( + image_group, + patch_size=output_kwargs["images_kwargs"].get("size"), + min_patches=min_patches, + max_patches=max_patches, + ) + images[index] = image_group + num_images = len(image_group) if (multi_page or crop_to_patches) else 1 + if box_single[0] is not None: + box_single = preprocess_box_annotation(box_single, image_group.size) + query = ( + f"{f'[{color_single}] ' if color_single is not None else ''}" + f"{str(box_single) if box_single[0] is not None else ''} " + "OCR" + f"{' with format' if format_output else ''}" + f"{' across multi pages' if multi_page else ''}" + f"{' upon the patch reference' if crop_to_patches else ''}" + ": " + ) + prompt = ( + self.message_start_token + + self.system_query + + self.message_end_token + + self.message_start_token + + "user\n" + + self.img_start_token + + self.img_pad_token * num_image_tokens * num_images + + self.img_end_token + + "\n" + + query + + self.message_end_token + + self.message_start_token + + "assistant\n" + ) + text.append(prompt) + elif crop_to_patches: + for index, (image_group, box_single, color_single) in enumerate(zip(images, box, color)): + image_group = self.image_processor.crop_image_to_patches( + image_group, + patch_size=output_kwargs["images_kwargs"].get("size"), + min_patches=min_patches, + max_patches=max_patches, + ) + images[index] = image_group + + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + if multi_page or crop_to_patches: + # flatten images + images = [image for image_group in images for image in image_group] + image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"]) + + return BatchFeature(data={**text_inputs, **image_inputs}) + + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please + refer to the docstring of this method for more information. + """ + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + """ + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + return list(tokenizer_input_names) + list(image_processor_input_names) + + +class GotOcr2MLPBlock(SamMLPBlock): + pass + + +class GotOcr2VisionAttention(SamVisionAttention): + pass + + +class GotOcr2VisionLayer(SamVisionLayer): + def __init__(self, config, window_size): + super().__init__() + self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.attn = GotOcr2VisionAttention(config, window_size) + self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.mlp = GotOcr2MLPBlock(config) + self.window_size = window_size + + +class GotOcr2VisionEncoder(SamVisionEncoder): + pass + + +class GotOcr2MultiModalProjector(nn.Module): + def __init__(self, config: GotOcr2Config): + super().__init__() + vision_output_channels = config.vision_config.output_channels + language_hidden_size = config.text_config.hidden_size + self.conv_upsampler1 = nn.Conv2d( + vision_output_channels, vision_output_channels * 2, kernel_size=3, stride=2, padding=1, bias=False + ) + self.conv_upsampler2 = nn.Conv2d( + vision_output_channels * 2, language_hidden_size, kernel_size=3, stride=2, padding=1, bias=False + ) + self.multimodal_projector = nn.Linear(language_hidden_size, language_hidden_size) + + def forward(self, vision_embeddings: torch.Tensor) -> torch.Tensor: + hidden_state = self.conv_upsampler1(vision_embeddings) + hidden_state = self.conv_upsampler2(hidden_state) + hidden_state = hidden_state.flatten(2).permute(0, 2, 1) + hidden_state = self.multimodal_projector(hidden_state) + return hidden_state + + +class GotOcr2CausalLMOutputWithPast(LlavaCausalLMOutputWithPast): + pass + + +class GotOcr2PreTrainedModel(LlavaPreTrainedModel): + pass + + +GOT_OCR2_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids) + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + pixel_values (`torch.FloatTensor` of shape `(seq_length, num_channels * image_size * image_size)): + The tensors corresponding to the input images. Pixel values can be obtained using + [`AutoImageProcessor`]. See [`GotOcr2ImageProcessor.__call__`] for details. [`GotOcr2Processor`] uses + [`GotOcr2ImageProcessor`] for processing images. +""" + + +class GotOcr2ForConditionalGeneration(LlavaForConditionalGeneration): + def __init__(self, config: GotOcr2Config): + super().__init__(config) + self.vision_tower = GotOcr2VisionEncoder(config.vision_config) + + self.multi_modal_projector = GotOcr2MultiModalProjector(config) + self.vocab_size = config.text_config.vocab_size + self.language_model = AutoModelForCausalLM.from_config(config.text_config) + + if self.language_model._tied_weights_keys is not None: + self._tied_weights_keys = [f"language_model.{k}" for k in self.language_model._tied_weights_keys] + + self.pad_token_id = config.pad_token_id + + self.post_init() + + def get_image_features( + self, + pixel_values: torch.FloatTensor, + ): + """ + Obtains image last hidden states from the vision tower and apply multimodal projection. + + Args: + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) + Returns: + image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + """ + image_outputs = self.vision_tower(pixel_values).last_hidden_state + return self.multi_modal_projector(image_outputs) + + @add_start_docstrings_to_model_forward(GOT_OCR2_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=GotOcr2CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + ) -> Union[Tuple, LlavaCausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + logits_to_keep (`int` or `torch.Tensor`, *optional*): + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. + This is useful when using packed tensor format (single dimension for batch and sequence length). + + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, GotOcr2ForConditionalGeneration, TextStreamer + + >>> model = GotOcr2ForConditionalGeneration.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf").to("cuda") + >>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + + >>> url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(image, return_tensors="pt", color="green").to("cuda") + + >>> # Generate + >>> streamer = TextStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True) + >>> generate_ids = model.generate( + ... **inputs, + ... do_sample=False, + ... tokenizer = processor.tokenizer, + ... stop_strings='<|im_end|>', + ... streamer=streamer, + ... max_new_tokens=4096, + ... ) + "You should keep in mind what features from the module should be used, especially + when you're planning to sell a template." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if pixel_values is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" + ) + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + image_features = self.get_image_features(pixel_values=pixel_values.to(inputs_embeds.dtype)) + n_image_tokens = (input_ids == self.config.image_token_index).sum() + n_image_features = image_features.shape[0] * image_features.shape[1] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1) + special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device) + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + logits_to_keep=logits_to_keep, + ) + + logits = outputs[0] + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + if attention_mask is not None: + # we use the input attention mask to shift the logits and labels, because it is 2D. + # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft + shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(logits.device) + shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous() + shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() + else: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1).to(shift_logits.device) + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return GotOcr2CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features if pixel_values is not None else None, + ) + + +__all__ = [ + "GotOcr2VisionConfig", + "GotOcr2Config", + "GotOcr2Processor", + "GotOcr2PreTrainedModel", + "GotOcr2ForConditionalGeneration", + "GotOcr2ImageProcessor", +] diff --git a/src/transformers/models/got_ocr2/processing_got_ocr2.py b/src/transformers/models/got_ocr2/processing_got_ocr2.py new file mode 100644 index 0000000000..636db765f9 --- /dev/null +++ b/src/transformers/models/got_ocr2/processing_got_ocr2.py @@ -0,0 +1,294 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/got_ocr2/modular_got_ocr2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_got_ocr2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2024 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. + + +from typing import List, Optional, Tuple, Union + +from transformers.processing_utils import ImagesKwargs, ProcessingKwargs, ProcessorMixin, TextKwargs, Unpack +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput + +from ...image_processing_utils import BatchFeature +from ...image_utils import ImageInput +from ...utils import is_vision_available, logging + + +if is_vision_available(): + from ...image_utils import load_images + +logger = logging.get_logger(__name__) + + +class GotOcr2TextKwargs(TextKwargs, total=False): + format: Optional[bool] + + +class GotOcr2ImagesKwargs(ImagesKwargs, total=False): + box: Optional[Union[List, Tuple[float, float], Tuple[float, float, float, float]]] + color: Optional[str] + num_image_tokens: Optional[int] + multi_page: Optional[bool] + crop_to_patches: Optional[bool] + min_patches: Optional[int] + max_patches: Optional[int] + + +class GotOcr2ProcessorKwargs(ProcessingKwargs, total=False): + text_kwargs: GotOcr2TextKwargs + images_kwargs: GotOcr2ImagesKwargs + _defaults = { + "text_kwargs": { + "padding": False, + "format": False, + }, + "images_kwargs": { + "num_image_tokens": 256, + "multi_page": False, + "crop_to_patches": False, + "min_patches": 1, + "max_patches": 12, + }, + } + + +def preprocess_box_annotation(box: Union[List, Tuple], image_size: Tuple[int, int]) -> List: + """ + Convert box annotation to the format [x1, y1, x2, y2] in the range [0, 1000]. + """ + width, height = image_size + if len(box) == 4: + box[0] = int(box[0] / width * 1000) + box[1] = int(box[1] / height * 1000) + box[2] = int(box[2] / width * 1000) + box[3] = int(box[3] / height * 1000) + else: + raise ValueError("Box must be a list or tuple of lists in the form [x1, y1, x2, y2].") + + return list(box) + + +class GotOcr2Processor(ProcessorMixin): + r""" + Constructs a GotOcr2 processor which wraps a [`GotOcr2ImageProcessor`] and + [`PretrainedTokenizerFast`] tokenizer into a single processor that inherits both the image processor and + tokenizer functionalities. See the [`~GotOcr2Processor.__call__`] and [`~GotOcr2Processor.decode`] for more information. + Args: + image_processor ([`GotOcr2ImageProcessor`], *optional*): + The image processor is a required input. + tokenizer ([`PreTrainedTokenizer`, `PreTrainedTokenizerFast`], *optional*): + The tokenizer is a required input. + chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages + in a chat into a tokenizable string. + """ + + attributes = ["image_processor", "tokenizer"] + valid_kwargs = ["chat_template"] + image_processor_class = "GotOcr2ImageProcessor" + tokenizer_class = "PreTrainedTokenizerFast" + + def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs): + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + self.message_start_token = "<|im_start|>" + self.message_end_token = "<|im_end|>" + self.img_start_token = "" + self.img_end_token = "" + self.img_pad_token = "" + self.system_query = "system\nYou should follow the instructions carefully and explain your answers in detail." + + def _make_list_of_inputs(self, images, text, box, color, multi_page): + if not isinstance(images, (list, tuple)): + images = [images] + if multi_page: + logger.warning("Multi-page inference is enabled but only one image is passed.") + images = [images] + elif isinstance(images[0], (list, tuple)) and not multi_page: + raise ValueError("Nested images are only supported with `multi_page` set to `True`.") + elif not isinstance(images[0], (list, tuple)) and multi_page: + images = [images] + + if isinstance(text, str): + text = [text] + + if not isinstance(box[0], (list, tuple)): + # Use the same box for all images + box = [box for _ in range(len(images))] + if not isinstance(color, (list, tuple)): + color = [color for _ in range(len(images))] + + return images, text, box, color + + def __call__( + self, + images: Optional[ImageInput] = None, + text: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None, + audio=None, + videos=None, + **kwargs: Unpack[GotOcr2ProcessorKwargs], + ) -> BatchFeature: + """ + Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text` + and `kwargs` arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizerFast.__call__`] to encode the text if `text` + is not `None`, otherwise encode default OCR queries which depends on the `format`, `box`, `color`, `multi_page` and + `crop_to_patches` arguments. To prepare the vision inputs, this method forwards the `images` and `kwrags` arguments to + GotOcr2ImageProcessor's [`~GotOcr2ImageProcessor.__call__`] if `images` is not `None`. + + Args: + images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`): + The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch + tensor. Both channels-first and channels-last formats are supported. + text (`str`, `List[str]`, `List[List[str]]`): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + format (`bool`, *optional*): + If set, will add the format token to the query, and the model will return the OCR result with formatting. + box (`List[float]`, `List[Tuple[float, float]]`, `List[Tuple[float, float, float, float]]`, *optional*): + The box annotation to be added to the query. If a list of floats or a tuple of floats is provided, it + will be interpreted as [x1, y1, x2, y2]. If a list of tuples is provided, each tuple should be in the + form (x1, y1, x2, y2). + color (`str`, *optional*): + The color annotation to be added to the query. The model will return the OCR result within the box with + the specified color. + multi_page (`bool`, *optional*): + If set, will enable multi-page inference. The model will return the OCR result across multiple pages. + crop_to_patches (`bool`, *optional*): + If set, will crop the image to patches. The model will return the OCR result upon the patch reference. + min_patches (`int`, *optional*): + The minimum number of patches to be cropped from the image. Only used when `crop_to_patches` is set to + `True`. + max_patches (`int`, *optional*): + The maximum number of patches to be cropped from the image. Only used when `crop_to_patches` is set to + `True`. + + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Acceptable values are: + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + - `'jax'`: Return JAX `jnp.ndarray` objects. + + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + + output_kwargs = self._merge_kwargs( + GotOcr2ProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + format_output = output_kwargs["text_kwargs"].pop("format") + num_image_tokens = output_kwargs["images_kwargs"].pop("num_image_tokens") + box = output_kwargs["images_kwargs"].pop("box", [None]) + color = output_kwargs["images_kwargs"].pop("color", None) + multi_page = output_kwargs["images_kwargs"].pop("multi_page") + crop_to_patches = output_kwargs["images_kwargs"].pop("crop_to_patches") + min_patches = output_kwargs["images_kwargs"].pop("min_patches") + max_patches = output_kwargs["images_kwargs"].pop("max_patches") + + images, text, box, color = self._make_list_of_inputs(images, text, box, color, multi_page) + + # Load images as we need to know the image size + images = load_images(images) + if text is None: + text = [] + for index, (image_group, box_single, color_single) in enumerate(zip(images, box, color)): + if crop_to_patches: + image_group = self.image_processor.crop_image_to_patches( + image_group, + patch_size=output_kwargs["images_kwargs"].get("size"), + min_patches=min_patches, + max_patches=max_patches, + ) + images[index] = image_group + num_images = len(image_group) if (multi_page or crop_to_patches) else 1 + if box_single[0] is not None: + box_single = preprocess_box_annotation(box_single, image_group.size) + query = ( + f"{f'[{color_single}] ' if color_single is not None else ''}" + f"{str(box_single) if box_single[0] is not None else ''} " + "OCR" + f"{' with format' if format_output else ''}" + f"{' across multi pages' if multi_page else ''}" + f"{' upon the patch reference' if crop_to_patches else ''}" + ": " + ) + prompt = ( + self.message_start_token + + self.system_query + + self.message_end_token + + self.message_start_token + + "user\n" + + self.img_start_token + + self.img_pad_token * num_image_tokens * num_images + + self.img_end_token + + "\n" + + query + + self.message_end_token + + self.message_start_token + + "assistant\n" + ) + text.append(prompt) + elif crop_to_patches: + for index, (image_group, box_single, color_single) in enumerate(zip(images, box, color)): + image_group = self.image_processor.crop_image_to_patches( + image_group, + patch_size=output_kwargs["images_kwargs"].get("size"), + min_patches=min_patches, + max_patches=max_patches, + ) + images[index] = image_group + + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + if multi_page or crop_to_patches: + # flatten images + images = [image for image_group in images for image in image_group] + image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"]) + + return BatchFeature(data={**text_inputs, **image_inputs}) + + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please + refer to the docstring of this method for more information. + """ + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + """ + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + return list(tokenizer_input_names) + list(image_processor_input_names) + + +__all__ = ["GotOcr2Processor"] diff --git a/src/transformers/models/llava/modeling_llava.py b/src/transformers/models/llava/modeling_llava.py index 67313c8f55..50b666758f 100644 --- a/src/transformers/models/llava/modeling_llava.py +++ b/src/transformers/models/llava/modeling_llava.py @@ -498,7 +498,7 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin): image_sizes=image_sizes, ) - n_image_tokens = (input_ids == self.config.image_token_index).sum().item() + n_image_tokens = (input_ids == self.config.image_token_index).sum() n_image_features = image_features.shape[0] * image_features.shape[1] if n_image_tokens != n_image_features: raise ValueError( diff --git a/src/transformers/utils/dummy_pt_objects.py b/src/transformers/utils/dummy_pt_objects.py index f379119289..dbfb1ef534 100644 --- a/src/transformers/utils/dummy_pt_objects.py +++ b/src/transformers/utils/dummy_pt_objects.py @@ -4637,6 +4637,20 @@ class GLPNPreTrainedModel(metaclass=DummyObject): requires_backends(self, ["torch"]) +class GotOcr2ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GotOcr2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + class GPT2DoubleHeadsModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/transformers/utils/dummy_vision_objects.py b/src/transformers/utils/dummy_vision_objects.py index 96eca67f39..f9802ba42b 100644 --- a/src/transformers/utils/dummy_vision_objects.py +++ b/src/transformers/utils/dummy_vision_objects.py @@ -289,6 +289,13 @@ class GLPNImageProcessor(metaclass=DummyObject): requires_backends(self, ["vision"]) +class GotOcr2ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + class GroundingDinoImageProcessor(metaclass=DummyObject): _backends = ["vision"] diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index 170128ce22..321803a217 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -1650,7 +1650,7 @@ class GenerationTesterMixin: # checks without adding test complexity. Ditto for `pixel_values_videos` and `pixel_values_images` pixel_values_is_mutually_exclusive = any( model_name in model_class.__name__.lower() - for model_name in ["llava", "idefics2", "idefics3", "mllama", "paligemma", "emu3"] + for model_name in ["llava", "idefics2", "idefics3", "mllama", "paligemma", "emu3", "gotocr2"] ) if pixel_values_is_mutually_exclusive: inputs_dict.pop("pixel_values", None) diff --git a/tests/models/got_ocr2/__init__.py b/tests/models/got_ocr2/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/models/got_ocr2/test_image_processing_got_ocr2.py b/tests/models/got_ocr2/test_image_processing_got_ocr2.py new file mode 100644 index 0000000000..c4e75feee6 --- /dev/null +++ b/tests/models/got_ocr2/test_image_processing_got_ocr2.py @@ -0,0 +1,115 @@ +# coding=utf-8 +# Copyright 2022 HuggingFace Inc. +# +# 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.testing_utils import require_torch, require_vision +from transformers.utils import is_vision_available + +from ...test_image_processing_common import ImageProcessingTestMixin, prepare_image_inputs + + +if is_vision_available(): + from transformers import GotOcr2ImageProcessor + + +class GotOcr2ImageProcessingTester(unittest.TestCase): + def __init__( + self, + parent, + batch_size=7, + num_channels=3, + image_size=18, + min_resolution=30, + max_resolution=400, + do_resize=True, + size=None, + do_normalize=True, + do_pad=False, + image_mean=[0.48145466, 0.4578275, 0.40821073], + image_std=[0.26862954, 0.26130258, 0.27577711], + do_convert_rgb=True, + ): + super().__init__() + size = size if size is not None else {"height": 20, "width": 20} + self.parent = parent + self.batch_size = batch_size + self.num_channels = num_channels + self.image_size = image_size + self.min_resolution = min_resolution + self.max_resolution = max_resolution + self.do_resize = do_resize + self.size = size + self.do_normalize = do_normalize + self.image_mean = image_mean + self.image_std = image_std + self.do_pad = do_pad + self.do_convert_rgb = do_convert_rgb + + def prepare_image_processor_dict(self): + return { + "do_resize": self.do_resize, + "size": self.size, + "do_normalize": self.do_normalize, + "image_mean": self.image_mean, + "image_std": self.image_std, + "do_convert_rgb": self.do_convert_rgb, + "do_pad": self.do_pad, + } + + def expected_output_image_shape(self, images): + return self.num_channels, self.size["height"], self.size["width"] + + def prepare_image_inputs(self, equal_resolution=False, numpify=False, torchify=False): + return prepare_image_inputs( + batch_size=self.batch_size, + num_channels=self.num_channels, + min_resolution=self.min_resolution, + max_resolution=self.max_resolution, + equal_resolution=equal_resolution, + numpify=numpify, + torchify=torchify, + ) + + +@require_torch +@require_vision +class GotOcr2ProcessingTest(ImageProcessingTestMixin, unittest.TestCase): + image_processing_class = GotOcr2ImageProcessor if is_vision_available() else None + + def setUp(self): + super().setUp() + self.image_processor_tester = GotOcr2ImageProcessingTester(self) + + @property + def image_processor_dict(self): + return self.image_processor_tester.prepare_image_processor_dict() + + def test_image_processor_properties(self): + image_processor = self.image_processing_class(**self.image_processor_dict) + self.assertTrue(hasattr(image_processor, "do_resize")) + self.assertTrue(hasattr(image_processor, "size")) + self.assertTrue(hasattr(image_processor, "do_normalize")) + self.assertTrue(hasattr(image_processor, "image_mean")) + self.assertTrue(hasattr(image_processor, "image_std")) + self.assertTrue(hasattr(image_processor, "do_convert_rgb")) + + def test_crop_to_patches(self): + image_processor = self.image_processing_class(**self.image_processor_dict) + image = self.image_processor_tester.prepare_image_inputs(equal_resolution=True)[0] + processed_images = image_processor.crop_image_to_patches(image, 1, 6, use_thumbnail=True) + self.assertEqual(len(processed_images), 5) + self.assertEqual(processed_images[0].size, (20, 20)) diff --git a/tests/models/got_ocr2/test_modeling_got_ocr2.py b/tests/models/got_ocr2/test_modeling_got_ocr2.py new file mode 100644 index 0000000000..ba3755237d --- /dev/null +++ b/tests/models/got_ocr2/test_modeling_got_ocr2.py @@ -0,0 +1,386 @@ +# coding=utf-8 +# Copyright 2024 The Qwen team, Alibaba Group 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 GotOcr2 model.""" + +import unittest + +from transformers import ( + AutoProcessor, + GotOcr2Config, + is_torch_available, + is_vision_available, +) +from transformers.testing_utils import cleanup, require_torch, slow, torch_device + +from ...generation.test_utils import GenerationTesterMixin +from ...test_configuration_common import ConfigTester +from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor +from ...test_pipeline_mixin import PipelineTesterMixin + + +if is_torch_available(): + import torch + + from transformers import ( + GotOcr2ForConditionalGeneration, + ) + + +if is_vision_available(): + from transformers.image_utils import load_image + + +class GotOcr2VisionText2TextModelTester: + def __init__( + self, + parent, + batch_size=3, + seq_length=7, + num_channels=3, + ignore_index=-100, + image_size=64, + bos_token_id=0, + eos_token_id=0, + pad_token_id=0, + image_token_index=1, + model_type="got_ocr2", + is_training=True, + text_config={ + "model_type": "qwen2", + "vocab_size": 99, + "hidden_size": 128, + "intermediate_size": 37, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "output_channels": 64, + "hidden_act": "silu", + "max_position_embeddings": 512, + "rope_theta": 10000, + "mlp_ratio": 4, + "tie_word_embeddings": True, + }, + vision_config={ + "num_hidden_layers": 2, + "output_channels": 64, + "hidden_act": "quick_gelu", + "hidden_size": 32, + "mlp_dim": 128, + "num_attention_heads": 4, + "patch_size": 2, + "image_size": 64, + }, + ): + self.parent = parent + self.ignore_index = ignore_index + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.image_token_index = image_token_index + self.model_type = model_type + self.text_config = text_config + self.vision_config = vision_config + self.batch_size = batch_size + self.num_channels = num_channels + self.image_size = image_size + self.is_training = is_training + self.num_image_tokens = 64 + self.seq_length = seq_length + self.num_image_tokens + + self.num_hidden_layers = text_config["num_hidden_layers"] + self.vocab_size = text_config["vocab_size"] + self.hidden_size = text_config["hidden_size"] + self.num_attention_heads = text_config["num_attention_heads"] + + def get_config(self): + return GotOcr2Config( + text_config=self.text_config, + vision_config=self.vision_config, + model_type=self.model_type, + bos_token_id=self.bos_token_id, + eos_token_id=self.eos_token_id, + pad_token_id=self.pad_token_id, + image_token_index=self.image_token_index, + ) + + def prepare_config_and_inputs(self): + config = self.get_config() + pixel_values = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]) + + return config, pixel_values + + def prepare_config_and_inputs_for_common(self): + config_and_inputs = self.prepare_config_and_inputs() + config, pixel_values = config_and_inputs + input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) + attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=torch_device) + + # input_ids[:, -1] = self.pad_token_id + input_ids[input_ids == self.image_token_index] = self.pad_token_id + input_ids[:, : self.num_image_tokens] = self.image_token_index + + inputs_dict = { + "pixel_values": pixel_values, + "input_ids": input_ids, + "attention_mask": attention_mask, + } + return config, inputs_dict + + def create_and_check_model_fp16_forward(self, config, input_ids, pixel_values, attention_mask): + model = GotOcr2ForConditionalGeneration(config=config) + model.to(torch_device) + model.half() + model.eval() + logits = model( + input_ids=input_ids, + attention_mask=attention_mask, + pixel_values=pixel_values.to(torch.bfloat16), + return_dict=True, + )["logits"] + self.parent.assertFalse(torch.isnan(logits).any().item()) + + def create_and_check_model_fp16_autocast_forward(self, config, input_ids, pixel_values, attention_mask): + config.torch_dtype = torch.float16 + model = GotOcr2ForConditionalGeneration(config=config) + model.to(torch_device) + model.eval() + with torch.autocast(device_type="cuda", dtype=torch.float16): + logits = model( + input_ids=input_ids, + attention_mask=attention_mask, + pixel_values=pixel_values.to(torch.bfloat16), + return_dict=True, + )["logits"] + self.parent.assertFalse(torch.isnan(logits).any().item()) + + +@require_torch +class GotOcr2ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase): + all_model_classes = (GotOcr2ForConditionalGeneration,) if is_torch_available() else () + all_generative_model_classes = (GotOcr2ForConditionalGeneration,) if is_torch_available() else () + pipeline_model_mapping = ( + { + "image-to-text": GotOcr2ForConditionalGeneration, + "image-text-to-text": GotOcr2ForConditionalGeneration, + } + if is_torch_available() + else {} + ) + test_headmasking = False + test_pruning = False + + def setUp(self): + self.model_tester = GotOcr2VisionText2TextModelTester(self) + self.config_tester = ConfigTester(self, config_class=GotOcr2Config, has_text_modality=False) + + def test_config(self): + self.config_tester.run_common_tests() + + def test_initialization(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + configs_no_init = _config_zero_init(config) + for model_class in self.all_model_classes: + model = model_class(config=configs_no_init) + for name, param in model.named_parameters(): + if param.requires_grad: + self.assertIn( + ((param.data.mean() * 1e9).round() / 1e9).item(), + [0.0, 1.0], + msg=f"Parameter {name} of model {model_class} seems not properly initialized", + ) + + # overwrite inputs_embeds tests because we need to delete "pixel values" for LVLMs + def test_inputs_embeds(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config) + model.to(torch_device) + model.eval() + + inputs = self._prepare_for_class(inputs_dict, model_class) + + input_ids = inputs["input_ids"] + del inputs["input_ids"] + del inputs["pixel_values"] + + wte = model.get_input_embeddings() + inputs["inputs_embeds"] = wte(input_ids) + + with torch.no_grad(): + model(**inputs) + + # overwrite inputs_embeds tests because we need to delete "pixel values" for LVLMs + # while some other models require pixel_values to be present + def test_inputs_embeds_matches_input_ids(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config) + model.to(torch_device) + model.eval() + + inputs = self._prepare_for_class(inputs_dict, model_class) + input_ids = inputs["input_ids"] + del inputs["input_ids"] + del inputs["pixel_values"] + + inputs_embeds = model.get_input_embeddings()(input_ids) + + with torch.no_grad(): + out_ids = model(input_ids=input_ids, **inputs)[0] + out_embeds = model(inputs_embeds=inputs_embeds, **inputs)[0] + torch.testing.assert_close(out_embeds, out_ids) + + @unittest.skip( + reason="VLMs can't generate from inputs embeds and pixels. This can be tested as part of bacbone LM, no need to run the test for VLMs" + ) + def test_generate_from_inputs_embeds_with_static_cache(self): + pass + + @unittest.skip( + reason="GotOcr2's language backbone is Qwen2 which uses GQA so the KV cache is a non standard format" + ) + def test_past_key_values_format(self): + pass + + @unittest.skip( + reason="GotOcr2 needs a dynamic control flow to pass pixel values to the forward function only in the first generation step" + ) + def test_generate_compile_1_end_to_end(self): + pass + + @unittest.skip("FlashAttention only support fp16 and bf16 data type") + def test_flash_attn_2_fp32_ln(self): + pass + + +@require_torch +class GotOcr2IntegrationTest(unittest.TestCase): + def setUp(self): + self.processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + + def tearDown(self): + cleanup(torch_device, gc_collect=True) + + @slow + def test_small_model_integration_test_got_ocr_stop_strings(self): + model_id = "stepfun-ai/GOT-OCR-2.0-hf" + model = GotOcr2ForConditionalGeneration.from_pretrained(model_id, device_map=torch_device) + image = load_image( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_ocr/resolve/main/iam_picture.jpeg" + ) + + inputs = self.processor(image, return_tensors="pt").to(torch_device) + generate_ids = model.generate( + **inputs, + do_sample=False, + num_beams=1, + tokenizer=self.processor.tokenizer, + stop_strings="<|im_end|>", + max_new_tokens=4096, + ) + decoded_output = self.processor.decode( + generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True + ) + expected_output = "industre" + self.assertEqual(decoded_output, expected_output) + + @slow + def test_small_model_integration_test_got_ocr_format(self): + model_id = "stepfun-ai/GOT-OCR-2.0-hf" + model = GotOcr2ForConditionalGeneration.from_pretrained(model_id, device_map=torch_device) + image = load_image( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" + ) + + inputs = self.processor(image, return_tensors="pt", format=True).to(torch_device) + generate_ids = model.generate(**inputs, do_sample=False, num_beams=1, max_new_tokens=4) + decoded_output = self.processor.decode( + generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True + ) + expected_output = "\\title{\nR" + self.assertEqual(decoded_output, expected_output) + + @slow + def test_small_model_integration_test_got_ocr_fine_grained(self): + model_id = "stepfun-ai/GOT-OCR-2.0-hf" + model = GotOcr2ForConditionalGeneration.from_pretrained(model_id, device_map=torch_device) + image = load_image( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" + ) + + inputs = self.processor(image, return_tensors="pt", color="green").to(torch_device) + generate_ids = model.generate(**inputs, do_sample=False, num_beams=1, max_new_tokens=4) + decoded_output = self.processor.decode( + generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True + ) + expected_output = "You should keep in" + self.assertEqual(decoded_output, expected_output) + + @slow + def test_small_model_integration_test_got_ocr_crop_to_patches(self): + model_id = "stepfun-ai/GOT-OCR-2.0-hf" + model = GotOcr2ForConditionalGeneration.from_pretrained(model_id, device_map=torch_device) + image = load_image( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/one_column.png" + ) + + inputs = self.processor(image, return_tensors="pt", crop_to_patches=True).to(torch_device) + generate_ids = model.generate(**inputs, do_sample=False, num_beams=1, max_new_tokens=4) + decoded_output = self.processor.decode( + generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True + ) + expected_output = "on developing architectural improvements" + self.assertEqual(decoded_output, expected_output) + + @slow + def test_small_model_integration_test_got_ocr_multi_pages(self): + model_id = "stepfun-ai/GOT-OCR-2.0-hf" + model = GotOcr2ForConditionalGeneration.from_pretrained(model_id, device_map=torch_device) + image1 = load_image( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/one_column.png" + ) + image2 = load_image( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" + ) + + inputs = self.processor([image1, image2], return_tensors="pt", multi_page=True).to(torch_device) + generate_ids = model.generate(**inputs, do_sample=False, num_beams=1, max_new_tokens=4) + decoded_output = self.processor.decode( + generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True + ) + expected_output = "on developing architectural improvements" + self.assertEqual(decoded_output, expected_output) + + @slow + def test_small_model_integration_test_got_ocr_batched(self): + model_id = "stepfun-ai/GOT-OCR-2.0-hf" + model = GotOcr2ForConditionalGeneration.from_pretrained(model_id, device_map=torch_device) + image1 = load_image( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" + ) + image2 = load_image( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" + ) + + inputs = self.processor([image1, image2], return_tensors="pt").to(torch_device) + generate_ids = model.generate(**inputs, do_sample=False, num_beams=1, max_new_tokens=4) + decoded_output = self.processor.batch_decode( + generate_ids[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True + ) + expected_output = ["Reducing the number", "R&D QUALITY"] + self.assertEqual(decoded_output, expected_output) diff --git a/tests/models/got_ocr2/test_processor_got_ocr2.py b/tests/models/got_ocr2/test_processor_got_ocr2.py new file mode 100644 index 0000000000..530d8234ce --- /dev/null +++ b/tests/models/got_ocr2/test_processor_got_ocr2.py @@ -0,0 +1,77 @@ +# Copyright 2024 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 shutil +import tempfile +import unittest + +from transformers import AutoProcessor, GotOcr2Processor, PreTrainedTokenizerFast +from transformers.testing_utils import require_vision +from transformers.utils import is_vision_available + +from ...test_processing_common import ProcessorTesterMixin + + +if is_vision_available(): + from transformers import GotOcr2ImageProcessor + + +@require_vision +class GotOcr2ProcessorTest(ProcessorTesterMixin, unittest.TestCase): + processor_class = GotOcr2Processor + + def setUp(self): + self.tmpdirname = tempfile.mkdtemp() + + image_processor = GotOcr2ImageProcessor() + tokenizer = PreTrainedTokenizerFast.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf") + processor_kwargs = self.prepare_processor_dict() + processor = GotOcr2Processor(image_processor, tokenizer, **processor_kwargs) + processor.save_pretrained(self.tmpdirname) + + def get_tokenizer(self, **kwargs): + return AutoProcessor.from_pretrained(self.tmpdirname, **kwargs).tokenizer + + def get_image_processor(self, **kwargs): + return AutoProcessor.from_pretrained(self.tmpdirname, **kwargs).image_processor + + def tearDown(self): + shutil.rmtree(self.tmpdirname) + + def test_ocr_queries(self): + processor = self.get_processor() + image_input = self.prepare_image_inputs() + inputs = processor(image_input, return_tensors="pt") + self.assertEqual(inputs["input_ids"].shape, (1, 286)) + self.assertEqual(inputs["pixel_values"].shape, (1, 3, 384, 384)) + + inputs = processor(image_input, return_tensors="pt", format=True) + self.assertEqual(inputs["input_ids"].shape, (1, 288)) + self.assertEqual(inputs["pixel_values"].shape, (1, 3, 384, 384)) + + inputs = processor(image_input, return_tensors="pt", color="red") + self.assertEqual(inputs["input_ids"].shape, (1, 290)) + self.assertEqual(inputs["pixel_values"].shape, (1, 3, 384, 384)) + + inputs = processor(image_input, return_tensors="pt", box=[0, 0, 100, 100]) + self.assertEqual(inputs["input_ids"].shape, (1, 303)) + self.assertEqual(inputs["pixel_values"].shape, (1, 3, 384, 384)) + + inputs = processor([image_input, image_input], return_tensors="pt", multi_page=True, format=True) + self.assertEqual(inputs["input_ids"].shape, (1, 547)) + self.assertEqual(inputs["pixel_values"].shape, (2, 3, 384, 384)) + + inputs = processor(image_input, return_tensors="pt", crop_to_patches=True, max_patches=6) + self.assertEqual(inputs["input_ids"].shape, (1, 1826)) + self.assertEqual(inputs["pixel_values"].shape, (7, 3, 384, 384)) diff --git a/tests/models/qwen2_vl/test_modeling_qwen2_vl.py b/tests/models/qwen2_vl/test_modeling_qwen2_vl.py index 8864185abf..05cf22a3fd 100644 --- a/tests/models/qwen2_vl/test_modeling_qwen2_vl.py +++ b/tests/models/qwen2_vl/test_modeling_qwen2_vl.py @@ -327,7 +327,7 @@ class Qwen2VLModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCas pass @unittest.skip( - reason="VLMs can't generate from inputs embeds and pixels. This can be tested as part of bacbone LM, no need to run the tes for VLMs" + reason="VLMs can't generate from inputs embeds and pixels. This can be tested as part of bacbone LM, no need to run the test for VLMs" ) def test_generate_from_inputs_embeds_with_static_cache(self): pass