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
@@ -872,6 +872,8 @@
|
||||
title: FLAVA
|
||||
- local: model_doc/git
|
||||
title: GIT
|
||||
- local: model_doc/got_ocr2
|
||||
title: GOT-OCR2
|
||||
- local: model_doc/grounding-dino
|
||||
title: Grounding DINO
|
||||
- local: model_doc/groupvit
|
||||
|
||||
@@ -161,6 +161,7 @@ Flax), PyTorch, and/or TensorFlow.
|
||||
| [GIT](model_doc/git) | ✅ | ❌ | ❌ |
|
||||
| [GLM](model_doc/glm) | ✅ | ❌ | ❌ |
|
||||
| [GLPN](model_doc/glpn) | ✅ | ❌ | ❌ |
|
||||
| [GOT-OCR2](model_doc/got_ocr2) | ✅ | ❌ | ❌ |
|
||||
| [GPT Neo](model_doc/gpt_neo) | ✅ | ❌ | ✅ |
|
||||
| [GPT NeoX](model_doc/gpt_neox) | ✅ | ❌ | ❌ |
|
||||
| [GPT NeoX Japanese](model_doc/gpt_neox_japanese) | ✅ | ❌ | ❌ |
|
||||
|
||||
269
docs/source/en/model_doc/got_ocr2.md
Normal file
269
docs/source/en/model_doc/got_ocr2.md
Normal 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 people’snusage due to the growing demand for intelligent processing of man-made opticalncharacters. In this paper, we collectively refer to all artificial optical signals (e.g., plain texts, math/molecular formulas, tables, charts, sheet music, and even geometric shapes) as "characters" and propose the General OCR Theory along with an excellent model, namely GOT, to promote the arrival of OCR-2.0. The GOT, with 580M parameters, is a unified, elegant, and end-to-end model, consisting of a high-compression encoder and a long-contexts decoder. As an OCR-2.0 model, GOT can handle all the above "characters" under various OCR tasks. On the input side, the model supports commonly used scene- and document-style images in slice and whole-page styles. On the output side, GOT can generate plain or formatted results (markdown/tikz/smiles/kern) via an easy prompt. Besides, the model enjoys interactive OCR features, i.e., region-level recognition guided by coordinates or colors. Furthermore, we also adapt dynamic resolution and multipage OCR technologies to GOT for better practicality. In experiments, we provide sufficient results to prove the superiority of our model.*
|
||||
|
||||
<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 you’re planning to sell a template."
|
||||
```
|
||||
|
||||
### Inference on general OCR data example: sheet music
|
||||
|
||||
Although this implementation of the model will only output plain text, the outputs can be further processed to render the desired format, with packages like `pdftex`, `mathpix`, `matplotlib`, `tikz`, `verovio` or `pyecharts`.
|
||||
Here is an example of how to process sheet music:
|
||||
|
||||
```python
|
||||
>>> from transformers import AutoProcessor, AutoModelForImageTextToText
|
||||
>>> import verovio
|
||||
|
||||
>>> device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
>>> model = AutoModelForImageTextToText.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf", device_map=device)
|
||||
>>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf")
|
||||
|
||||
>>> image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/sheet_music.png"
|
||||
>>> inputs = processor(image, return_tensors="pt", format=True).to(device)
|
||||
|
||||
>>> generate_ids = model.generate(
|
||||
... **inputs,
|
||||
... do_sample=False,
|
||||
... tokenizer=processor.tokenizer,
|
||||
... stop_strings="<|im_end|>",
|
||||
... max_new_tokens=4096,
|
||||
... )
|
||||
|
||||
>>> outputs = processor.decode(generate_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
|
||||
>>> tk = verovio.toolkit()
|
||||
>>> tk.loadData(outputs)
|
||||
>>> tk.setOptions(
|
||||
... {
|
||||
... "pageWidth": 2100,
|
||||
... "pageHeight": 800,
|
||||
... "footer": "none",
|
||||
... "barLineWidth": 0.5,
|
||||
... "beamMaxSlope": 15,
|
||||
... "staffLineWidth": 0.2,
|
||||
... "spacingStaff": 6,
|
||||
... }
|
||||
... )
|
||||
>>> tk.getPageCount()
|
||||
>>> svg = tk.renderToSVG()
|
||||
>>> svg = svg.replace('overflow="inherit"', 'overflow="visible"')
|
||||
>>> with open("output.svg", "w") as f:
|
||||
>>> f.write(svg)
|
||||
```
|
||||
<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
|
||||
|
||||
@@ -52,6 +52,7 @@ FlashAttention-2 is currently supported for the following architectures:
|
||||
* [Emu3](https://huggingface.co/docs/transformers/model_doc/emu3)
|
||||
* [Gemma](https://huggingface.co/docs/transformers/model_doc/gemma#transformers.GemmaModel)
|
||||
* [Gemma2](https://huggingface.co/docs/transformers/model_doc/gemma2#transformers.Gemma2Model)
|
||||
* [GotOcr2](https://huggingface.co/docs/transformers/model_doc/got_ocr2#transformers.GotOcr2ForConditionalGeneration)
|
||||
* [GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2)
|
||||
* [GPTBigCode](https://huggingface.co/docs/transformers/model_doc/gpt_bigcode#transformers.GPTBigCodeModel)
|
||||
* [GPTNeo](https://huggingface.co/docs/transformers/model_doc/gpt_neo#transformers.GPTNeoModel)
|
||||
@@ -253,6 +254,7 @@ For now, Transformers supports SDPA inference and training for the following arc
|
||||
* [Falcon](https://huggingface.co/docs/transformers/model_doc/falcon#transformers.FalconModel)
|
||||
* [Gemma](https://huggingface.co/docs/transformers/model_doc/gemma#transformers.GemmaModel)
|
||||
* [Gemma2](https://huggingface.co/docs/transformers/model_doc/gemma2#transformers.Gemma2Model)
|
||||
* [GotOcr2](https://huggingface.co/docs/transformers/model_doc/got_ocr2#transformers.GotOcr2ForConditionalGeneration)
|
||||
* [Granite](https://huggingface.co/docs/transformers/model_doc/granite#transformers.GraniteModel)
|
||||
* [GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2)
|
||||
* [GPTBigCode](https://huggingface.co/docs/transformers/model_doc/gpt_bigcode#transformers.GPTBigCodeModel)
|
||||
|
||||
@@ -476,6 +476,11 @@ _import_structure = {
|
||||
],
|
||||
"models.glm": ["GlmConfig"],
|
||||
"models.glpn": ["GLPNConfig"],
|
||||
"models.got_ocr2": [
|
||||
"GotOcr2Config",
|
||||
"GotOcr2Processor",
|
||||
"GotOcr2VisionConfig",
|
||||
],
|
||||
"models.gpt2": [
|
||||
"GPT2Config",
|
||||
"GPT2Tokenizer",
|
||||
@@ -1238,6 +1243,7 @@ else:
|
||||
_import_structure["models.flava"].extend(["FlavaFeatureExtractor", "FlavaImageProcessor", "FlavaProcessor"])
|
||||
_import_structure["models.fuyu"].extend(["FuyuImageProcessor", "FuyuProcessor"])
|
||||
_import_structure["models.glpn"].extend(["GLPNFeatureExtractor", "GLPNImageProcessor"])
|
||||
_import_structure["models.got_ocr2"].extend(["GotOcr2ImageProcessor"])
|
||||
_import_structure["models.grounding_dino"].extend(["GroundingDinoImageProcessor"])
|
||||
_import_structure["models.idefics"].extend(["IdeficsImageProcessor"])
|
||||
_import_structure["models.idefics2"].extend(["Idefics2ImageProcessor"])
|
||||
@@ -2426,6 +2432,12 @@ else:
|
||||
"GLPNPreTrainedModel",
|
||||
]
|
||||
)
|
||||
_import_structure["models.got_ocr2"].extend(
|
||||
[
|
||||
"GotOcr2ForConditionalGeneration",
|
||||
"GotOcr2PreTrainedModel",
|
||||
]
|
||||
)
|
||||
_import_structure["models.gpt2"].extend(
|
||||
[
|
||||
"GPT2DoubleHeadsModel",
|
||||
@@ -5540,6 +5552,7 @@ if TYPE_CHECKING:
|
||||
)
|
||||
from .models.glm import GlmConfig
|
||||
from .models.glpn import GLPNConfig
|
||||
from .models.got_ocr2 import GotOcr2Config, GotOcr2Processor, GotOcr2VisionConfig
|
||||
from .models.gpt2 import (
|
||||
GPT2Config,
|
||||
GPT2Tokenizer,
|
||||
@@ -6342,6 +6355,7 @@ if TYPE_CHECKING:
|
||||
)
|
||||
from .models.fuyu import FuyuImageProcessor, FuyuProcessor
|
||||
from .models.glpn import GLPNFeatureExtractor, GLPNImageProcessor
|
||||
from .models.got_ocr2 import GotOcr2ImageProcessor
|
||||
from .models.grounding_dino import GroundingDinoImageProcessor
|
||||
from .models.idefics import IdeficsImageProcessor
|
||||
from .models.idefics2 import Idefics2ImageProcessor
|
||||
@@ -7346,6 +7360,10 @@ if TYPE_CHECKING:
|
||||
GLPNModel,
|
||||
GLPNPreTrainedModel,
|
||||
)
|
||||
from .models.got_ocr2 import (
|
||||
GotOcr2ForConditionalGeneration,
|
||||
GotOcr2PreTrainedModel,
|
||||
)
|
||||
from .models.gpt2 import (
|
||||
GPT2DoubleHeadsModel,
|
||||
GPT2ForQuestionAnswering,
|
||||
|
||||
@@ -106,6 +106,7 @@ from . import (
|
||||
git,
|
||||
glm,
|
||||
glpn,
|
||||
got_ocr2,
|
||||
gpt2,
|
||||
gpt_bigcode,
|
||||
gpt_neo,
|
||||
|
||||
@@ -124,6 +124,7 @@ CONFIG_MAPPING_NAMES = OrderedDict(
|
||||
("git", "GitConfig"),
|
||||
("glm", "GlmConfig"),
|
||||
("glpn", "GLPNConfig"),
|
||||
("got_ocr2", "GotOcr2Config"),
|
||||
("gpt-sw3", "GPT2Config"),
|
||||
("gpt2", "GPT2Config"),
|
||||
("gpt_bigcode", "GPTBigCodeConfig"),
|
||||
@@ -450,6 +451,7 @@ MODEL_NAMES_MAPPING = OrderedDict(
|
||||
("git", "GIT"),
|
||||
("glm", "GLM"),
|
||||
("glpn", "GLPN"),
|
||||
("got_ocr2", "GOT-OCR2"),
|
||||
("gpt-sw3", "GPT-Sw3"),
|
||||
("gpt2", "OpenAI GPT-2"),
|
||||
("gpt_bigcode", "GPTBigCode"),
|
||||
|
||||
@@ -87,6 +87,7 @@ else:
|
||||
("fuyu", ("FuyuImageProcessor",)),
|
||||
("git", ("CLIPImageProcessor",)),
|
||||
("glpn", ("GLPNImageProcessor",)),
|
||||
("got_ocr2", ("GotOcr2ImageProcessor",)),
|
||||
("grounding-dino", ("GroundingDinoImageProcessor",)),
|
||||
("groupvit", ("CLIPImageProcessor",)),
|
||||
("hiera", ("BitImageProcessor",)),
|
||||
|
||||
@@ -119,6 +119,7 @@ MODEL_MAPPING_NAMES = OrderedDict(
|
||||
("git", "GitModel"),
|
||||
("glm", "GlmModel"),
|
||||
("glpn", "GLPNModel"),
|
||||
("got_ocr2", "GotOcr2ForConditionalGeneration"),
|
||||
("gpt-sw3", "GPT2Model"),
|
||||
("gpt2", "GPT2Model"),
|
||||
("gpt_bigcode", "GPTBigCodeModel"),
|
||||
@@ -512,6 +513,7 @@ MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
|
||||
("gemma2", "Gemma2ForCausalLM"),
|
||||
("git", "GitForCausalLM"),
|
||||
("glm", "GlmForCausalLM"),
|
||||
("got_ocr2", "GotOcr2ForConditionalGeneration"),
|
||||
("gpt-sw3", "GPT2LMHeadModel"),
|
||||
("gpt2", "GPT2LMHeadModel"),
|
||||
("gpt_bigcode", "GPTBigCodeForCausalLM"),
|
||||
@@ -811,6 +813,7 @@ MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES = OrderedDict(
|
||||
("emu3", "Emu3ForConditionalGeneration"),
|
||||
("fuyu", "FuyuForCausalLM"),
|
||||
("git", "GitForCausalLM"),
|
||||
("got_ocr2", "GotOcr2ForConditionalGeneration"),
|
||||
("idefics", "IdeficsForVisionText2Text"),
|
||||
("idefics2", "Idefics2ForConditionalGeneration"),
|
||||
("idefics3", "Idefics3ForConditionalGeneration"),
|
||||
|
||||
@@ -63,6 +63,7 @@ PROCESSOR_MAPPING_NAMES = OrderedDict(
|
||||
("flava", "FlavaProcessor"),
|
||||
("fuyu", "FuyuProcessor"),
|
||||
("git", "GitProcessor"),
|
||||
("got_ocr2", "GotOcr2Processor"),
|
||||
("grounding-dino", "GroundingDinoProcessor"),
|
||||
("groupvit", "CLIPProcessor"),
|
||||
("hubert", "Wav2Vec2Processor"),
|
||||
|
||||
31
src/transformers/models/got_ocr2/__init__.py
Normal file
31
src/transformers/models/got_ocr2/__init__.py
Normal 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__)
|
||||
212
src/transformers/models/got_ocr2/configuration_got_ocr2.py
Normal file
212
src/transformers/models/got_ocr2/configuration_got_ocr2.py
Normal 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"]
|
||||
@@ -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()
|
||||
482
src/transformers/models/got_ocr2/image_processing_got_ocr2.py
Normal file
482
src/transformers/models/got_ocr2/image_processing_got_ocr2.py
Normal 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"]
|
||||
1005
src/transformers/models/got_ocr2/modeling_got_ocr2.py
Normal file
1005
src/transformers/models/got_ocr2/modeling_got_ocr2.py
Normal file
File diff suppressed because it is too large
Load Diff
984
src/transformers/models/got_ocr2/modular_got_ocr2.py
Normal file
984
src/transformers/models/got_ocr2/modular_got_ocr2.py
Normal 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",
|
||||
]
|
||||
294
src/transformers/models/got_ocr2/processing_got_ocr2.py
Normal file
294
src/transformers/models/got_ocr2/processing_got_ocr2.py
Normal 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"]
|
||||
@@ -498,7 +498,7 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin):
|
||||
image_sizes=image_sizes,
|
||||
)
|
||||
|
||||
n_image_tokens = (input_ids == self.config.image_token_index).sum().item()
|
||||
n_image_tokens = (input_ids == self.config.image_token_index).sum()
|
||||
n_image_features = image_features.shape[0] * image_features.shape[1]
|
||||
if n_image_tokens != n_image_features:
|
||||
raise ValueError(
|
||||
|
||||
@@ -4637,6 +4637,20 @@ class GLPNPreTrainedModel(metaclass=DummyObject):
|
||||
requires_backends(self, ["torch"])
|
||||
|
||||
|
||||
class GotOcr2ForConditionalGeneration(metaclass=DummyObject):
|
||||
_backends = ["torch"]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
requires_backends(self, ["torch"])
|
||||
|
||||
|
||||
class GotOcr2PreTrainedModel(metaclass=DummyObject):
|
||||
_backends = ["torch"]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
requires_backends(self, ["torch"])
|
||||
|
||||
|
||||
class GPT2DoubleHeadsModel(metaclass=DummyObject):
|
||||
_backends = ["torch"]
|
||||
|
||||
|
||||
@@ -289,6 +289,13 @@ class GLPNImageProcessor(metaclass=DummyObject):
|
||||
requires_backends(self, ["vision"])
|
||||
|
||||
|
||||
class GotOcr2ImageProcessor(metaclass=DummyObject):
|
||||
_backends = ["vision"]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
requires_backends(self, ["vision"])
|
||||
|
||||
|
||||
class GroundingDinoImageProcessor(metaclass=DummyObject):
|
||||
_backends = ["vision"]
|
||||
|
||||
|
||||
@@ -1650,7 +1650,7 @@ class GenerationTesterMixin:
|
||||
# checks without adding test complexity. Ditto for `pixel_values_videos` and `pixel_values_images`
|
||||
pixel_values_is_mutually_exclusive = any(
|
||||
model_name in model_class.__name__.lower()
|
||||
for model_name in ["llava", "idefics2", "idefics3", "mllama", "paligemma", "emu3"]
|
||||
for model_name in ["llava", "idefics2", "idefics3", "mllama", "paligemma", "emu3", "gotocr2"]
|
||||
)
|
||||
if pixel_values_is_mutually_exclusive:
|
||||
inputs_dict.pop("pixel_values", None)
|
||||
|
||||
0
tests/models/got_ocr2/__init__.py
Normal file
0
tests/models/got_ocr2/__init__.py
Normal file
115
tests/models/got_ocr2/test_image_processing_got_ocr2.py
Normal file
115
tests/models/got_ocr2/test_image_processing_got_ocr2.py
Normal 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))
|
||||
386
tests/models/got_ocr2/test_modeling_got_ocr2.py
Normal file
386
tests/models/got_ocr2/test_modeling_got_ocr2.py
Normal 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)
|
||||
77
tests/models/got_ocr2/test_processor_got_ocr2.py
Normal file
77
tests/models/got_ocr2/test_processor_got_ocr2.py
Normal 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))
|
||||
@@ -327,7 +327,7 @@ class Qwen2VLModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCas
|
||||
pass
|
||||
|
||||
@unittest.skip(
|
||||
reason="VLMs can't generate from inputs embeds and pixels. This can be tested as part of bacbone LM, no need to run the tes for VLMs"
|
||||
reason="VLMs can't generate from inputs embeds and pixels. This can be tested as part of bacbone LM, no need to run the test for VLMs"
|
||||
)
|
||||
def test_generate_from_inputs_embeds_with_static_cache(self):
|
||||
pass
|
||||
|
||||
Reference in New Issue
Block a user