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
This commit is contained in:
Yoni Gozlan
2025-01-31 11:28:13 -05:00
committed by GitHub
parent 5bbee12ac9
commit 2b46943195
26 changed files with 4184 additions and 3 deletions

View File

@@ -872,6 +872,8 @@
title: FLAVA title: FLAVA
- local: model_doc/git - local: model_doc/git
title: GIT title: GIT
- local: model_doc/got_ocr2
title: GOT-OCR2
- local: model_doc/grounding-dino - local: model_doc/grounding-dino
title: Grounding DINO title: Grounding DINO
- local: model_doc/groupvit - local: model_doc/groupvit

View File

@@ -161,6 +161,7 @@ Flax), PyTorch, and/or TensorFlow.
| [GIT](model_doc/git) | ✅ | ❌ | ❌ | | [GIT](model_doc/git) | ✅ | ❌ | ❌ |
| [GLM](model_doc/glm) | ✅ | ❌ | ❌ | | [GLM](model_doc/glm) | ✅ | ❌ | ❌ |
| [GLPN](model_doc/glpn) | ✅ | ❌ | ❌ | | [GLPN](model_doc/glpn) | ✅ | ❌ | ❌ |
| [GOT-OCR2](model_doc/got_ocr2) | ✅ | ❌ | ❌ |
| [GPT Neo](model_doc/gpt_neo) | ✅ | ❌ | ✅ | | [GPT Neo](model_doc/gpt_neo) | ✅ | ❌ | ✅ |
| [GPT NeoX](model_doc/gpt_neox) | ✅ | ❌ | ❌ | | [GPT NeoX](model_doc/gpt_neox) | ✅ | ❌ | ❌ |
| [GPT NeoX Japanese](model_doc/gpt_neox_japanese) | ✅ | ❌ | ❌ | | [GPT NeoX Japanese](model_doc/gpt_neox_japanese) | ✅ | ❌ | ❌ |

View File

@@ -0,0 +1,269 @@
<!--Copyright 2024 StepFun and 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.
⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
rendered properly in your Markdown viewer.
-->
# 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 peoplesnusage 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.*
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/got_ocr_overview.png"
alt="drawing" width="600"/>
<small> GOT-OCR2 training stages. Taken from the <a href="https://arxiv.org/abs/2409.01704">original paper.</a> </small>
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 youre 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)
```
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/sheet_music.svg"
alt="drawing" width="600"/>
## GotOcr2Config
[[autodoc]] GotOcr2Config
## GotOcr2VisionConfig
[[autodoc]] GotOcr2VisionConfig
## GotOcr2ImageProcessor
[[autodoc]] GotOcr2ImageProcessor
## GotOcr2Processor
[[autodoc]] GotOcr2Processor
## GotOcr2ForConditionalGeneration
[[autodoc]] GotOcr2ForConditionalGeneration
- forward

View File

@@ -52,6 +52,7 @@ FlashAttention-2 is currently supported for the following architectures:
* [Emu3](https://huggingface.co/docs/transformers/model_doc/emu3) * [Emu3](https://huggingface.co/docs/transformers/model_doc/emu3)
* [Gemma](https://huggingface.co/docs/transformers/model_doc/gemma#transformers.GemmaModel) * [Gemma](https://huggingface.co/docs/transformers/model_doc/gemma#transformers.GemmaModel)
* [Gemma2](https://huggingface.co/docs/transformers/model_doc/gemma2#transformers.Gemma2Model) * [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) * [GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2)
* [GPTBigCode](https://huggingface.co/docs/transformers/model_doc/gpt_bigcode#transformers.GPTBigCodeModel) * [GPTBigCode](https://huggingface.co/docs/transformers/model_doc/gpt_bigcode#transformers.GPTBigCodeModel)
* [GPTNeo](https://huggingface.co/docs/transformers/model_doc/gpt_neo#transformers.GPTNeoModel) * [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) * [Falcon](https://huggingface.co/docs/transformers/model_doc/falcon#transformers.FalconModel)
* [Gemma](https://huggingface.co/docs/transformers/model_doc/gemma#transformers.GemmaModel) * [Gemma](https://huggingface.co/docs/transformers/model_doc/gemma#transformers.GemmaModel)
* [Gemma2](https://huggingface.co/docs/transformers/model_doc/gemma2#transformers.Gemma2Model) * [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) * [Granite](https://huggingface.co/docs/transformers/model_doc/granite#transformers.GraniteModel)
* [GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2) * [GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2)
* [GPTBigCode](https://huggingface.co/docs/transformers/model_doc/gpt_bigcode#transformers.GPTBigCodeModel) * [GPTBigCode](https://huggingface.co/docs/transformers/model_doc/gpt_bigcode#transformers.GPTBigCodeModel)

View File

@@ -476,6 +476,11 @@ _import_structure = {
], ],
"models.glm": ["GlmConfig"], "models.glm": ["GlmConfig"],
"models.glpn": ["GLPNConfig"], "models.glpn": ["GLPNConfig"],
"models.got_ocr2": [
"GotOcr2Config",
"GotOcr2Processor",
"GotOcr2VisionConfig",
],
"models.gpt2": [ "models.gpt2": [
"GPT2Config", "GPT2Config",
"GPT2Tokenizer", "GPT2Tokenizer",
@@ -1238,6 +1243,7 @@ else:
_import_structure["models.flava"].extend(["FlavaFeatureExtractor", "FlavaImageProcessor", "FlavaProcessor"]) _import_structure["models.flava"].extend(["FlavaFeatureExtractor", "FlavaImageProcessor", "FlavaProcessor"])
_import_structure["models.fuyu"].extend(["FuyuImageProcessor", "FuyuProcessor"]) _import_structure["models.fuyu"].extend(["FuyuImageProcessor", "FuyuProcessor"])
_import_structure["models.glpn"].extend(["GLPNFeatureExtractor", "GLPNImageProcessor"]) _import_structure["models.glpn"].extend(["GLPNFeatureExtractor", "GLPNImageProcessor"])
_import_structure["models.got_ocr2"].extend(["GotOcr2ImageProcessor"])
_import_structure["models.grounding_dino"].extend(["GroundingDinoImageProcessor"]) _import_structure["models.grounding_dino"].extend(["GroundingDinoImageProcessor"])
_import_structure["models.idefics"].extend(["IdeficsImageProcessor"]) _import_structure["models.idefics"].extend(["IdeficsImageProcessor"])
_import_structure["models.idefics2"].extend(["Idefics2ImageProcessor"]) _import_structure["models.idefics2"].extend(["Idefics2ImageProcessor"])
@@ -2426,6 +2432,12 @@ else:
"GLPNPreTrainedModel", "GLPNPreTrainedModel",
] ]
) )
_import_structure["models.got_ocr2"].extend(
[
"GotOcr2ForConditionalGeneration",
"GotOcr2PreTrainedModel",
]
)
_import_structure["models.gpt2"].extend( _import_structure["models.gpt2"].extend(
[ [
"GPT2DoubleHeadsModel", "GPT2DoubleHeadsModel",
@@ -5540,6 +5552,7 @@ if TYPE_CHECKING:
) )
from .models.glm import GlmConfig from .models.glm import GlmConfig
from .models.glpn import GLPNConfig from .models.glpn import GLPNConfig
from .models.got_ocr2 import GotOcr2Config, GotOcr2Processor, GotOcr2VisionConfig
from .models.gpt2 import ( from .models.gpt2 import (
GPT2Config, GPT2Config,
GPT2Tokenizer, GPT2Tokenizer,
@@ -6342,6 +6355,7 @@ if TYPE_CHECKING:
) )
from .models.fuyu import FuyuImageProcessor, FuyuProcessor from .models.fuyu import FuyuImageProcessor, FuyuProcessor
from .models.glpn import GLPNFeatureExtractor, GLPNImageProcessor from .models.glpn import GLPNFeatureExtractor, GLPNImageProcessor
from .models.got_ocr2 import GotOcr2ImageProcessor
from .models.grounding_dino import GroundingDinoImageProcessor from .models.grounding_dino import GroundingDinoImageProcessor
from .models.idefics import IdeficsImageProcessor from .models.idefics import IdeficsImageProcessor
from .models.idefics2 import Idefics2ImageProcessor from .models.idefics2 import Idefics2ImageProcessor
@@ -7346,6 +7360,10 @@ if TYPE_CHECKING:
GLPNModel, GLPNModel,
GLPNPreTrainedModel, GLPNPreTrainedModel,
) )
from .models.got_ocr2 import (
GotOcr2ForConditionalGeneration,
GotOcr2PreTrainedModel,
)
from .models.gpt2 import ( from .models.gpt2 import (
GPT2DoubleHeadsModel, GPT2DoubleHeadsModel,
GPT2ForQuestionAnswering, GPT2ForQuestionAnswering,

View File

@@ -106,6 +106,7 @@ from . import (
git, git,
glm, glm,
glpn, glpn,
got_ocr2,
gpt2, gpt2,
gpt_bigcode, gpt_bigcode,
gpt_neo, gpt_neo,

View File

@@ -124,6 +124,7 @@ CONFIG_MAPPING_NAMES = OrderedDict(
("git", "GitConfig"), ("git", "GitConfig"),
("glm", "GlmConfig"), ("glm", "GlmConfig"),
("glpn", "GLPNConfig"), ("glpn", "GLPNConfig"),
("got_ocr2", "GotOcr2Config"),
("gpt-sw3", "GPT2Config"), ("gpt-sw3", "GPT2Config"),
("gpt2", "GPT2Config"), ("gpt2", "GPT2Config"),
("gpt_bigcode", "GPTBigCodeConfig"), ("gpt_bigcode", "GPTBigCodeConfig"),
@@ -450,6 +451,7 @@ MODEL_NAMES_MAPPING = OrderedDict(
("git", "GIT"), ("git", "GIT"),
("glm", "GLM"), ("glm", "GLM"),
("glpn", "GLPN"), ("glpn", "GLPN"),
("got_ocr2", "GOT-OCR2"),
("gpt-sw3", "GPT-Sw3"), ("gpt-sw3", "GPT-Sw3"),
("gpt2", "OpenAI GPT-2"), ("gpt2", "OpenAI GPT-2"),
("gpt_bigcode", "GPTBigCode"), ("gpt_bigcode", "GPTBigCode"),

View File

@@ -87,6 +87,7 @@ else:
("fuyu", ("FuyuImageProcessor",)), ("fuyu", ("FuyuImageProcessor",)),
("git", ("CLIPImageProcessor",)), ("git", ("CLIPImageProcessor",)),
("glpn", ("GLPNImageProcessor",)), ("glpn", ("GLPNImageProcessor",)),
("got_ocr2", ("GotOcr2ImageProcessor",)),
("grounding-dino", ("GroundingDinoImageProcessor",)), ("grounding-dino", ("GroundingDinoImageProcessor",)),
("groupvit", ("CLIPImageProcessor",)), ("groupvit", ("CLIPImageProcessor",)),
("hiera", ("BitImageProcessor",)), ("hiera", ("BitImageProcessor",)),

View File

@@ -119,6 +119,7 @@ MODEL_MAPPING_NAMES = OrderedDict(
("git", "GitModel"), ("git", "GitModel"),
("glm", "GlmModel"), ("glm", "GlmModel"),
("glpn", "GLPNModel"), ("glpn", "GLPNModel"),
("got_ocr2", "GotOcr2ForConditionalGeneration"),
("gpt-sw3", "GPT2Model"), ("gpt-sw3", "GPT2Model"),
("gpt2", "GPT2Model"), ("gpt2", "GPT2Model"),
("gpt_bigcode", "GPTBigCodeModel"), ("gpt_bigcode", "GPTBigCodeModel"),
@@ -512,6 +513,7 @@ MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
("gemma2", "Gemma2ForCausalLM"), ("gemma2", "Gemma2ForCausalLM"),
("git", "GitForCausalLM"), ("git", "GitForCausalLM"),
("glm", "GlmForCausalLM"), ("glm", "GlmForCausalLM"),
("got_ocr2", "GotOcr2ForConditionalGeneration"),
("gpt-sw3", "GPT2LMHeadModel"), ("gpt-sw3", "GPT2LMHeadModel"),
("gpt2", "GPT2LMHeadModel"), ("gpt2", "GPT2LMHeadModel"),
("gpt_bigcode", "GPTBigCodeForCausalLM"), ("gpt_bigcode", "GPTBigCodeForCausalLM"),
@@ -811,6 +813,7 @@ MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES = OrderedDict(
("emu3", "Emu3ForConditionalGeneration"), ("emu3", "Emu3ForConditionalGeneration"),
("fuyu", "FuyuForCausalLM"), ("fuyu", "FuyuForCausalLM"),
("git", "GitForCausalLM"), ("git", "GitForCausalLM"),
("got_ocr2", "GotOcr2ForConditionalGeneration"),
("idefics", "IdeficsForVisionText2Text"), ("idefics", "IdeficsForVisionText2Text"),
("idefics2", "Idefics2ForConditionalGeneration"), ("idefics2", "Idefics2ForConditionalGeneration"),
("idefics3", "Idefics3ForConditionalGeneration"), ("idefics3", "Idefics3ForConditionalGeneration"),

View File

@@ -63,6 +63,7 @@ PROCESSOR_MAPPING_NAMES = OrderedDict(
("flava", "FlavaProcessor"), ("flava", "FlavaProcessor"),
("fuyu", "FuyuProcessor"), ("fuyu", "FuyuProcessor"),
("git", "GitProcessor"), ("git", "GitProcessor"),
("got_ocr2", "GotOcr2Processor"),
("grounding-dino", "GroundingDinoProcessor"), ("grounding-dino", "GroundingDinoProcessor"),
("groupvit", "CLIPProcessor"), ("groupvit", "CLIPProcessor"),
("hubert", "Wav2Vec2Processor"), ("hubert", "Wav2Vec2Processor"),

View File

@@ -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__)

View File

@@ -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"]

View File

@@ -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)]
+ [
"<ref>",
"</ref>",
"<box>",
"</box>",
"<quad>",
"</quad>",
"<img>",
"</img>",
"<imgpad>",
]
)
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()

View File

@@ -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"]

File diff suppressed because it is too large Load Diff

View File

@@ -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 = "<img>"
self.img_end_token = "</img>"
self.img_pad_token = "<imgpad>"
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",
]

View File

@@ -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 = "<img>"
self.img_end_token = "</img>"
self.img_pad_token = "<imgpad>"
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"]

View File

@@ -498,7 +498,7 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin):
image_sizes=image_sizes, 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] n_image_features = image_features.shape[0] * image_features.shape[1]
if n_image_tokens != n_image_features: if n_image_tokens != n_image_features:
raise ValueError( raise ValueError(

View File

@@ -4637,6 +4637,20 @@ class GLPNPreTrainedModel(metaclass=DummyObject):
requires_backends(self, ["torch"]) 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): class GPT2DoubleHeadsModel(metaclass=DummyObject):
_backends = ["torch"] _backends = ["torch"]

View File

@@ -289,6 +289,13 @@ class GLPNImageProcessor(metaclass=DummyObject):
requires_backends(self, ["vision"]) requires_backends(self, ["vision"])
class GotOcr2ImageProcessor(metaclass=DummyObject):
_backends = ["vision"]
def __init__(self, *args, **kwargs):
requires_backends(self, ["vision"])
class GroundingDinoImageProcessor(metaclass=DummyObject): class GroundingDinoImageProcessor(metaclass=DummyObject):
_backends = ["vision"] _backends = ["vision"]

View File

@@ -1650,7 +1650,7 @@ class GenerationTesterMixin:
# checks without adding test complexity. Ditto for `pixel_values_videos` and `pixel_values_images` # checks without adding test complexity. Ditto for `pixel_values_videos` and `pixel_values_images`
pixel_values_is_mutually_exclusive = any( pixel_values_is_mutually_exclusive = any(
model_name in model_class.__name__.lower() 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: if pixel_values_is_mutually_exclusive:
inputs_dict.pop("pixel_values", None) inputs_dict.pop("pixel_values", None)

View File

View File

@@ -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))

View File

@@ -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)

View File

@@ -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))

View File

@@ -327,7 +327,7 @@ class Qwen2VLModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCas
pass pass
@unittest.skip( @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): def test_generate_from_inputs_embeds_with_static_cache(self):
pass pass