From 0f68a7f4088cac3a97b28e339d8b8d72c0edf93b Mon Sep 17 00:00:00 2001 From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> Date: Wed, 22 Mar 2023 16:53:52 +0100 Subject: [PATCH] Add Pix2Struct (#21400) * v1 all keys match * clean up * forward pass ok * add correct image transform * generate works, logits matching * clean up * more refactor * revert * revert * clean up * clean ups * clean up * refactor * refactor * fix doc * fix tokenizer test * fix toctree * revert toctree * oops * few fixes * replace to `pixel_embeds` * make fixup * test processing & feat extractor * fix some tests * more fixes * make fixup * clean up * more clean up * add a single slow test * fix test * make fixup * fix * fix authors * fix toctree * update docs * add docstring * revert change * Update src/transformers/models/pix2struct/__init__.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * fix tokenizer * fix processor test * fix test * make fixup * refactor * fix config * Update src/transformers/models/pix2struct/image_processing_pix2struct.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * format * fix * Update src/transformers/models/pix2struct/image_processing_pix2struct.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * make fixup * add docstring * fix issues * fix * fix * fix * add slow test * fix * fix * fix batched issue * fix training issues * fix ci test * fix slow test * fix conversion script * remove unneeded classes * fix slow test * fix require backends * fix masked fill * revert * fix softmax * add large models support * fix conditional generation * few fixes * add instructions * rm unneeded file * Update src/transformers/models/pix2struct/convert_pix2struct_original_pytorch_to_hf.py * fix ci test * fix ci test really * Apply suggestions from code review Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * fix nit * fix nits * fix image processors nits * docstring * clean up * fix nit * fix tests * docstring nit * fix reshape * Update src/transformers/models/pix2struct/image_processing_pix2struct.py Co-authored-by: NielsRogge <48327001+NielsRogge@users.noreply.github.com> * fix nit * fix repetition * refactor processor * make patch size consistent * refactor forward * fix docstring * fix max_patches issue * update docstirng * update docstring * fix coped from * add skip reasons * few fixes * Update src/transformers/models/pix2struct/image_processing_pix2struct.py Co-authored-by: NielsRogge <48327001+NielsRogge@users.noreply.github.com> * format * fix doctests * refactor and fix * fix doc build issue * fix processor test * small fix conversion script * replace correct weights * make fixup * fix some issues * Apply suggestions from code review Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * revert config and fixes * Update src/transformers/models/pix2struct/image_processing_pix2struct.py Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * more details * fixes * fix processor * fix processor test * fix * Apply suggestions from code review Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * make fixup * fix processor * Update src/transformers/models/pix2struct/modeling_pix2struct.py Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * add copied * make fixup * fix copies * update docstring * refactor * fix docstring * fix conversion script * fix vqa issue * replace to `flattened_patches` * nit * fix numpy issue * fix image processors * add batched vqa support * fix vqa conversion * make fixup * fix conversion script * Apply suggestions from code review Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * make fixup * add correct docstring * update docstring * fix module level + channel dim * use `make_list_of_images` * refactor * correct docstring * fix authors * remove `data_format` * add header text test * Apply suggestions from code review Co-authored-by: NielsRogge <48327001+NielsRogge@users.noreply.github.com> Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * make fixup * add checkpoints --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> Co-authored-by: NielsRogge <48327001+NielsRogge@users.noreply.github.com> --- README.md | 1 + README_es.md | 1 + README_hd.md | 1 + README_ja.md | 1 + README_ko.md | 1 + README_zh-hans.md | 1 + README_zh-hant.md | 1 + docs/source/en/_toctree.yml | 2 + docs/source/en/index.mdx | 2 + docs/source/en/model_doc/pix2struct.mdx | 73 + src/transformers/__init__.py | 32 + src/transformers/models/__init__.py | 1 + .../models/auto/configuration_auto.py | 3 + .../models/auto/image_processing_auto.py | 1 + .../models/auto/processing_auto.py | 1 + .../models/auto/tokenization_auto.py | 1 + .../models/pix2struct/__init__.py | 86 + .../pix2struct/configuration_pix2struct.py | 412 ++++ ...nvert_pix2struct_original_pytorch_to_hf.py | 155 ++ .../pix2struct/image_processing_pix2struct.py | 418 ++++ .../models/pix2struct/modeling_pix2struct.py | 1823 +++++++++++++++++ .../pix2struct/processing_pix2struct.py | 162 ++ src/transformers/utils/dummy_pt_objects.py | 31 + .../utils/dummy_vision_objects.py | 7 + tests/models/pix2struct/__init__.py | 0 .../test_image_processing_pix2struct.py | 289 +++ .../pix2struct/test_modeling_pix2struct.py | 711 +++++++ .../pix2struct/test_processor_pix2struct.py | 192 ++ utils/check_repo.py | 3 + 29 files changed, 4412 insertions(+) create mode 100644 docs/source/en/model_doc/pix2struct.mdx create mode 100644 src/transformers/models/pix2struct/__init__.py create mode 100644 src/transformers/models/pix2struct/configuration_pix2struct.py create mode 100644 src/transformers/models/pix2struct/convert_pix2struct_original_pytorch_to_hf.py create mode 100644 src/transformers/models/pix2struct/image_processing_pix2struct.py create mode 100644 src/transformers/models/pix2struct/modeling_pix2struct.py create mode 100644 src/transformers/models/pix2struct/processing_pix2struct.py create mode 100644 tests/models/pix2struct/__init__.py create mode 100644 tests/models/pix2struct/test_image_processing_pix2struct.py create mode 100644 tests/models/pix2struct/test_modeling_pix2struct.py create mode 100644 tests/models/pix2struct/test_processor_pix2struct.py diff --git a/README.md b/README.md index 6397d14ada..209d3e21b2 100644 --- a/README.md +++ b/README.md @@ -399,6 +399,7 @@ Current number of checkpoints: ![](https://img.shields.io/endpoint?url=https://h 1. **[PEGASUS-X](https://huggingface.co/docs/transformers/model_doc/pegasus_x)** (from Google) released with the paper [Investigating Efficiently Extending Transformers for Long Input Summarization](https://arxiv.org/abs/2208.04347) by Jason Phang, Yao Zhao, and Peter J. Liu. 1. **[Perceiver IO](https://huggingface.co/docs/transformers/model_doc/perceiver)** (from Deepmind) released with the paper [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://arxiv.org/abs/2107.14795) by Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, João Carreira. 1. **[PhoBERT](https://huggingface.co/docs/transformers/model_doc/phobert)** (from VinAI Research) released with the paper [PhoBERT: Pre-trained language models for Vietnamese](https://www.aclweb.org/anthology/2020.findings-emnlp.92/) by Dat Quoc Nguyen and Anh Tuan Nguyen. +1. **[Pix2Struct](https://huggingface.co/docs/transformers/main/model_doc/pix2struct)** (from Google) released with the paper [Pix2Struct: Screenshot Parsing as Pretraining for Visual Language Understanding](https://arxiv.org/abs/2210.03347) by Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, Kristina Toutanova. 1. **[PLBart](https://huggingface.co/docs/transformers/model_doc/plbart)** (from UCLA NLP) released with the paper [Unified Pre-training for Program Understanding and Generation](https://arxiv.org/abs/2103.06333) by Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, Kai-Wei Chang. 1. **[PoolFormer](https://huggingface.co/docs/transformers/model_doc/poolformer)** (from Sea AI Labs) released with the paper [MetaFormer is Actually What You Need for Vision](https://arxiv.org/abs/2111.11418) by Yu, Weihao and Luo, Mi and Zhou, Pan and Si, Chenyang and Zhou, Yichen and Wang, Xinchao and Feng, Jiashi and Yan, Shuicheng. 1. **[ProphetNet](https://huggingface.co/docs/transformers/model_doc/prophetnet)** (from Microsoft Research) released with the paper [ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training](https://arxiv.org/abs/2001.04063) by Yu Yan, Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang and Ming Zhou. diff --git a/README_es.md b/README_es.md index 899140210c..b350d45d40 100644 --- a/README_es.md +++ b/README_es.md @@ -387,6 +387,7 @@ Número actual de puntos de control: ![](https://img.shields.io/endpoint?url=htt 1. **[PEGASUS-X](https://huggingface.co/docs/transformers/model_doc/pegasus_x)** (from Google) released with the paper [Investigating Efficiently Extending Transformers for Long Input Summarization](https://arxiv.org/abs/2208.04347) by Jason Phang, Yao Zhao, and Peter J. Liu. 1. **[Perceiver IO](https://huggingface.co/docs/transformers/model_doc/perceiver)** (from Deepmind) released with the paper [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://arxiv.org/abs/2107.14795) by Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, João Carreira. 1. **[PhoBERT](https://huggingface.co/docs/transformers/model_doc/phobert)** (from VinAI Research) released with the paper [PhoBERT: Pre-trained language models for Vietnamese](https://www.aclweb.org/anthology/2020.findings-emnlp.92/) by Dat Quoc Nguyen and Anh Tuan Nguyen. +1. **[Pix2Struct](https://huggingface.co/docs/transformers/main/model_doc/pix2struct)** (from Google) released with the paper [Pix2Struct: Screenshot Parsing as Pretraining for Visual Language Understanding](https://arxiv.org/abs/2210.03347) by Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, Kristina Toutanova. 1. **[PLBart](https://huggingface.co/docs/transformers/model_doc/plbart)** (from UCLA NLP) released with the paper [Unified Pre-training for Program Understanding and Generation](https://arxiv.org/abs/2103.06333) by Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, Kai-Wei Chang. 1. **[PoolFormer](https://huggingface.co/docs/transformers/model_doc/poolformer)** (from Sea AI Labs) released with the paper [MetaFormer is Actually What You Need for Vision](https://arxiv.org/abs/2111.11418) by Yu, Weihao and Luo, Mi and Zhou, Pan and Si, Chenyang and Zhou, Yichen and Wang, Xinchao and Feng, Jiashi and Yan, Shuicheng. 1. **[ProphetNet](https://huggingface.co/docs/transformers/model_doc/prophetnet)** (from Microsoft Research) released with the paper [ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training](https://arxiv.org/abs/2001.04063) by Yu Yan, Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang and Ming Zhou. diff --git a/README_hd.md b/README_hd.md index 826306cf67..71bf9b61b5 100644 --- a/README_hd.md +++ b/README_hd.md @@ -359,6 +359,7 @@ conda install -c huggingface transformers 1. **[PEGASUS-X](https://huggingface.co/docs/transformers/model_doc/pegasus_x)** (Google की ओर से) साथ में दिया गया पेपर [लंबे इनपुट सारांश के लिए ट्रांसफ़ॉर्मरों को बेहतर तरीके से एक्सटेंड करना](https://arxiv .org/abs/2208.04347) जेसन फांग, याओ झाओ, पीटर जे लियू द्वारा। 1. **[Perceiver IO](https://huggingface.co/docs/transformers/model_doc/perceiver)** (दीपमाइंड से) साथ में पेपर [पर्सीवर आईओ: संरचित इनपुट और आउटपुट के लिए एक सामान्य वास्तुकला] (https://arxiv.org/abs/2107.14795) एंड्रयू जेगल, सेबेस्टियन बोरग्यूड, जीन-बैप्टिस्ट अलायराक, कार्ल डोर्श, कैटलिन इओनेस्कु, डेविड द्वारा डिंग, स्कंद कोप्पुला, डैनियल ज़ोरान, एंड्रयू ब्रॉक, इवान शेलहैमर, ओलिवियर हेनाफ, मैथ्यू एम। बोट्विनिक, एंड्रयू ज़िसरमैन, ओरिओल विनियल्स, जोआओ कैरेरा द्वारा पोस्ट किया गया। 1. **[PhoBERT](https://huggingface.co/docs/transformers/model_doc/phobert)** (VinAI Research से) कागज के साथ [PhoBERT: वियतनामी के लिए पूर्व-प्रशिक्षित भाषा मॉडल](https://www .aclweb.org/anthology/2020.findings-emnlp.92/) डैट क्वोक गुयेन और अन्ह तुआन गुयेन द्वारा पोस्ट किया गया। +1. **[Pix2Struct](https://huggingface.co/docs/transformers/main/model_doc/pix2struct)** (Google से) Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, Kristina Toutanova. द्वाराअनुसंधान पत्र [Pix2Struct: Screenshot Parsing as Pretraining for Visual Language Understanding](https://arxiv.org/abs/2210.03347) के साथ जारी किया गया 1. **[PLBart](https://huggingface.co/docs/transformers/model_doc/plbart)** (UCLA NLP से) साथ वाला पेपर [प्रोग्राम अंडरस्टैंडिंग एंड जेनरेशन के लिए यूनिफाइड प्री-ट्रेनिंग](https://arxiv .org/abs/2103.06333) वसी उद्दीन अहमद, सैकत चक्रवर्ती, बैशाखी रे, काई-वेई चांग द्वारा। 1. **[PoolFormer](https://huggingface.co/docs/transformers/model_doc/poolformer)** (from Sea AI Labs) released with the paper [MetaFormer is Actually What You Need for Vision](https://arxiv.org/abs/2111.11418) by Yu, Weihao and Luo, Mi and Zhou, Pan and Si, Chenyang and Zhou, Yichen and Wang, Xinchao and Feng, Jiashi and Yan, Shuicheng. 1. **[ProphetNet](https://huggingface.co/docs/transformers/model_doc/prophetnet)** (माइक्रोसॉफ्ट रिसर्च से) साथ में पेपर [ProphetNet: प्रेडिक्टिंग फ्यूचर एन-ग्राम फॉर सीक्वेंस-टू-सीक्वेंस प्री-ट्रेनिंग ](https://arxiv.org/abs/2001.04063) यू यान, वीज़ेन क्यूई, येयुन गोंग, दयाहेंग लियू, नान डुआन, जिउशेंग चेन, रुओफ़ेई झांग और मिंग झोउ द्वारा पोस्ट किया गया। diff --git a/README_ja.md b/README_ja.md index b45cc68ea6..75f45c4462 100644 --- a/README_ja.md +++ b/README_ja.md @@ -421,6 +421,7 @@ Flax、PyTorch、TensorFlowをcondaでインストールする方法は、それ 1. **[PEGASUS-X](https://huggingface.co/docs/transformers/model_doc/pegasus_x)** (Google から) Jason Phang, Yao Zhao, and Peter J. Liu から公開された研究論文: [Investigating Efficiently Extending Transformers for Long Input Summarization](https://arxiv.org/abs/2208.04347) 1. **[Perceiver IO](https://huggingface.co/docs/transformers/model_doc/perceiver)** (Deepmind から) Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, João Carreira から公開された研究論文: [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://arxiv.org/abs/2107.14795) 1. **[PhoBERT](https://huggingface.co/docs/transformers/model_doc/phobert)** (VinAI Research から) Dat Quoc Nguyen and Anh Tuan Nguyen から公開された研究論文: [PhoBERT: Pre-trained language models for Vietnamese](https://www.aclweb.org/anthology/2020.findings-emnlp.92/) +1. **[Pix2Struct](https://huggingface.co/docs/transformers/main/model_doc/pix2struct)** (Google から) Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, Kristina Toutanova. から公開された研究論文 [Pix2Struct: Screenshot Parsing as Pretraining for Visual Language Understanding](https://arxiv.org/abs/2210.03347) 1. **[PLBart](https://huggingface.co/docs/transformers/model_doc/plbart)** (UCLA NLP から) Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, Kai-Wei Chang から公開された研究論文: [Unified Pre-training for Program Understanding and Generation](https://arxiv.org/abs/2103.06333) 1. **[PoolFormer](https://huggingface.co/docs/transformers/model_doc/poolformer)** (Sea AI Labs から) Yu, Weihao and Luo, Mi and Zhou, Pan and Si, Chenyang and Zhou, Yichen and Wang, Xinchao and Feng, Jiashi and Yan, Shuicheng から公開された研究論文: [MetaFormer is Actually What You Need for Vision](https://arxiv.org/abs/2111.11418) 1. **[ProphetNet](https://huggingface.co/docs/transformers/model_doc/prophetnet)** (Microsoft Research から) Yu Yan, Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang and Ming Zhou から公開された研究論文: [ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training](https://arxiv.org/abs/2001.04063) diff --git a/README_ko.md b/README_ko.md index a5c0b8cf1e..703c898707 100644 --- a/README_ko.md +++ b/README_ko.md @@ -336,6 +336,7 @@ Flax, PyTorch, TensorFlow 설치 페이지에서 이들을 conda로 설치하는 1. **[PEGASUS-X](https://huggingface.co/docs/transformers/model_doc/pegasus_x)** (Google 에서) Jason Phang, Yao Zhao, Peter J. Liu 의 [Investigating Efficiently Extending Transformers for Long Input Summarization](https://arxiv.org/abs/2208.04347) 논문과 함께 발표했습니다. 1. **[Perceiver IO](https://huggingface.co/docs/transformers/model_doc/perceiver)** (Deepmind 에서) Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, João Carreira 의 [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://arxiv.org/abs/2107.14795) 논문과 함께 발표했습니다. 1. **[PhoBERT](https://huggingface.co/docs/transformers/model_doc/phobert)** (VinAI Research 에서) Dat Quoc Nguyen and Anh Tuan Nguyen 의 [PhoBERT: Pre-trained language models for Vietnamese](https://www.aclweb.org/anthology/2020.findings-emnlp.92/) 논문과 함께 발표했습니다. +1. **[Pix2Struct](https://huggingface.co/docs/transformers/main/model_doc/pix2struct)** (Google 에서 제공)은 Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, Kristina Toutanova.의 [Pix2Struct: Screenshot Parsing as Pretraining for Visual Language Understanding](https://arxiv.org/abs/2210.03347)논문과 함께 발표했습니다. 1. **[PLBart](https://huggingface.co/docs/transformers/model_doc/plbart)** (UCLA NLP 에서) Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, Kai-Wei Chang 의 [Unified Pre-training for Program Understanding and Generation](https://arxiv.org/abs/2103.06333) 논문과 함께 발표했습니다. 1. **[PoolFormer](https://huggingface.co/docs/transformers/model_doc/poolformer)** (Sea AI Labs 에서) Yu, Weihao and Luo, Mi and Zhou, Pan and Si, Chenyang and Zhou, Yichen and Wang, Xinchao and Feng, Jiashi and Yan, Shuicheng 의 [MetaFormer is Actually What You Need for Vision](https://arxiv.org/abs/2111.11418) 논문과 함께 발표했습니다. 1. **[ProphetNet](https://huggingface.co/docs/transformers/model_doc/prophetnet)** (Microsoft Research 에서) Yu Yan, Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang and Ming Zhou 의 [ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training](https://arxiv.org/abs/2001.04063) 논문과 함께 발표했습니다. diff --git a/README_zh-hans.md b/README_zh-hans.md index 9ae3bc2449..8fa869b498 100644 --- a/README_zh-hans.md +++ b/README_zh-hans.md @@ -360,6 +360,7 @@ conda install -c huggingface transformers 1. **[PEGASUS-X](https://huggingface.co/docs/transformers/model_doc/pegasus_x)** (来自 Google) 伴随论文 [Investigating Efficiently Extending Transformers for Long Input Summarization](https://arxiv.org/abs/2208.04347) 由 Jason Phang, Yao Zhao, Peter J. Liu 发布。 1. **[Perceiver IO](https://huggingface.co/docs/transformers/model_doc/perceiver)** (来自 Deepmind) 伴随论文 [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://arxiv.org/abs/2107.14795) 由 Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, João Carreira 发布。 1. **[PhoBERT](https://huggingface.co/docs/transformers/model_doc/phobert)** (来自 VinAI Research) 伴随论文 [PhoBERT: Pre-trained language models for Vietnamese](https://www.aclweb.org/anthology/2020.findings-emnlp.92/) 由 Dat Quoc Nguyen and Anh Tuan Nguyen 发布。 +1. **[Pix2Struct](https://huggingface.co/docs/transformers/main/model_doc/pix2struct)** (来自 Google) 伴随论文 [Pix2Struct: Screenshot Parsing as Pretraining for Visual Language Understanding](https://arxiv.org/abs/2210.03347) 由 Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, Kristina Toutanova 发布。 1. **[PLBart](https://huggingface.co/docs/transformers/model_doc/plbart)** (来自 UCLA NLP) 伴随论文 [Unified Pre-training for Program Understanding and Generation](https://arxiv.org/abs/2103.06333) 由 Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, Kai-Wei Chang 发布。 1. **[PoolFormer](https://huggingface.co/docs/transformers/model_doc/poolformer)** (来自 Sea AI Labs) 伴随论文 [MetaFormer is Actually What You Need for Vision](https://arxiv.org/abs/2111.11418) 由 Yu, Weihao and Luo, Mi and Zhou, Pan and Si, Chenyang and Zhou, Yichen and Wang, Xinchao and Feng, Jiashi and Yan, Shuicheng 发布。 1. **[ProphetNet](https://huggingface.co/docs/transformers/model_doc/prophetnet)** (来自 Microsoft Research) 伴随论文 [ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training](https://arxiv.org/abs/2001.04063) 由 Yu Yan, Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang and Ming Zhou 发布。 diff --git a/README_zh-hant.md b/README_zh-hant.md index 53847bf673..eae33654a9 100644 --- a/README_zh-hant.md +++ b/README_zh-hant.md @@ -372,6 +372,7 @@ conda install -c huggingface transformers 1. **[PEGASUS-X](https://huggingface.co/docs/transformers/model_doc/pegasus_x)** (from Google) released with the paper [Investigating Efficiently Extending Transformers for Long Input Summarization](https://arxiv.org/abs/2208.04347) by Jason Phang, Yao Zhao, Peter J. Liu. 1. **[Perceiver IO](https://huggingface.co/docs/transformers/model_doc/perceiver)** (from Deepmind) released with the paper [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://arxiv.org/abs/2107.14795) by Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, João Carreira. 1. **[PhoBERT](https://huggingface.co/docs/transformers/model_doc/phobert)** (from VinAI Research) released with the paper [PhoBERT: Pre-trained language models for Vietnamese](https://www.aclweb.org/anthology/2020.findings-emnlp.92/) by Dat Quoc Nguyen and Anh Tuan Nguyen. +1. **[Pix2Struct](https://huggingface.co/docs/transformers/main/model_doc/pix2struct)** (from Google) released with the paper [Pix2Struct: Screenshot Parsing as Pretraining for Visual Language Understanding](https://arxiv.org/abs/2210.03347) by Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, Kristina Toutanova. 1. **[PLBart](https://huggingface.co/docs/transformers/model_doc/plbart)** (from UCLA NLP) released with the paper [Unified Pre-training for Program Understanding and Generation](https://arxiv.org/abs/2103.06333) by Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, Kai-Wei Chang. 1. **[PoolFormer](https://huggingface.co/docs/transformers/model_doc/poolformer)** (from Sea AI Labs) released with the paper [MetaFormer is Actually What You Need for Vision](https://arxiv.org/abs/2111.11418) by Yu, Weihao and Luo, Mi and Zhou, Pan and Si, Chenyang and Zhou, Yichen and Wang, Xinchao and Feng, Jiashi and Yan, Shuicheng. 1. **[ProphetNet](https://huggingface.co/docs/transformers/model_doc/prophetnet)** (from Microsoft Research) released with the paper [ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training](https://arxiv.org/abs/2001.04063) by Yu Yan, Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang and Ming Zhou. diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 7586985b11..c2e6933d44 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -596,6 +596,8 @@ title: OWL-ViT - local: model_doc/perceiver title: Perceiver + - local: model_doc/pix2struct + title: Pix2Struct - local: model_doc/speech-encoder-decoder title: Speech Encoder Decoder Models - local: model_doc/tapas diff --git a/docs/source/en/index.mdx b/docs/source/en/index.mdx index ea1ab27e79..eee2750345 100644 --- a/docs/source/en/index.mdx +++ b/docs/source/en/index.mdx @@ -173,6 +173,7 @@ The documentation is organized into five sections: 1. **[PEGASUS-X](model_doc/pegasus_x)** (from Google) released with the paper [Investigating Efficiently Extending Transformers for Long Input Summarization](https://arxiv.org/abs/2208.04347) by Jason Phang, Yao Zhao, and Peter J. Liu. 1. **[Perceiver IO](model_doc/perceiver)** (from Deepmind) released with the paper [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://arxiv.org/abs/2107.14795) by Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, João Carreira. 1. **[PhoBERT](model_doc/phobert)** (from VinAI Research) released with the paper [PhoBERT: Pre-trained language models for Vietnamese](https://www.aclweb.org/anthology/2020.findings-emnlp.92/) by Dat Quoc Nguyen and Anh Tuan Nguyen. +1. **[Pix2Struct](model_doc/pix2struct)** (from Google) released with the paper [Pix2Struct: Screenshot Parsing as Pretraining for Visual Language Understanding](https://arxiv.org/abs/2210.03347) by Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, Kristina Toutanova. 1. **[PLBart](model_doc/plbart)** (from UCLA NLP) released with the paper [Unified Pre-training for Program Understanding and Generation](https://arxiv.org/abs/2103.06333) by Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, Kai-Wei Chang. 1. **[PoolFormer](model_doc/poolformer)** (from Sea AI Labs) released with the paper [MetaFormer is Actually What You Need for Vision](https://arxiv.org/abs/2111.11418) by Yu, Weihao and Luo, Mi and Zhou, Pan and Si, Chenyang and Zhou, Yichen and Wang, Xinchao and Feng, Jiashi and Yan, Shuicheng. 1. **[ProphetNet](model_doc/prophetnet)** (from Microsoft Research) released with the paper [ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training](https://arxiv.org/abs/2001.04063) by Yu Yan, Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang and Ming Zhou. @@ -364,6 +365,7 @@ Flax), PyTorch, and/or TensorFlow. | Pegasus | ✅ | ✅ | ✅ | ✅ | ✅ | | PEGASUS-X | ❌ | ❌ | ✅ | ❌ | ❌ | | Perceiver | ✅ | ❌ | ✅ | ❌ | ❌ | +| Pix2Struct | ❌ | ❌ | ✅ | ❌ | ❌ | | PLBart | ✅ | ❌ | ✅ | ❌ | ❌ | | PoolFormer | ❌ | ❌ | ✅ | ❌ | ❌ | | ProphetNet | ✅ | ❌ | ✅ | ❌ | ❌ | diff --git a/docs/source/en/model_doc/pix2struct.mdx b/docs/source/en/model_doc/pix2struct.mdx new file mode 100644 index 0000000000..340b06c69b --- /dev/null +++ b/docs/source/en/model_doc/pix2struct.mdx @@ -0,0 +1,73 @@ + + +# Pix2Struct + +## Overview + +The Pix2Struct model was proposed in [Pix2Struct: Screenshot Parsing as Pretraining for Visual Language Understanding](https://arxiv.org/abs/2210.03347) by Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, Kristina Toutanova. + +The abstract from the paper is the following: + +> Visually-situated language is ubiquitous -- sources range from textbooks with diagrams to web pages with images and tables, to mobile apps with buttons and forms. Perhaps due to this diversity, previous work has typically relied on domain-specific recipes with limited sharing of the underlying data, model architectures, and objectives. We present Pix2Struct, a pretrained image-to-text model for purely visual language understanding, which can be finetuned on tasks containing visually-situated language. Pix2Struct is pretrained by learning to parse masked screenshots of web pages into simplified HTML. The web, with its richness of visual elements cleanly reflected in the HTML structure, provides a large source of pretraining data well suited to the diversity of downstream tasks. Intuitively, this objective subsumes common pretraining signals such as OCR, language modeling, image captioning. In addition to the novel pretraining strategy, we introduce a variable-resolution input representation and a more flexible integration of language and vision inputs, where language prompts such as questions are rendered directly on top of the input image. For the first time, we show that a single pretrained model can achieve state-of-the-art results in six out of nine tasks across four domains: documents, illustrations, user interfaces, and natural images. + +Tips: + +Pix2Struct has been fine tuned on a variety of tasks and datasets, ranging from image captioning, visual question answering (VQA) over different inputs (books, charts, science diagrams), captioning UI components etc. The full list can be found in Table 1 of the paper. +We therefore advise you to use these models for the tasks they have been fine tuned on. For instance, if you want to use Pix2Struct for UI captioning, you should use the model fine tuned on the UI dataset. If you want to use Pix2Struct for image captioning, you should use the model fine tuned on the natural images captioning dataset and so on. + +This model was contributed by [ybelkada](https://huggingface.co/ybelkada). +The original code can be found [here](https://github.com/google-research/pix2struct). + +## Resources: + +- [Paper](https://arxiv.org/abs/2210.03347) +- [Fine-tuning Notebook](https://github.com/huggingface/notebooks/blob/main/examples/image_captioning_pix2struct.ipynb) +- [All models](https://huggingface.co/models?search=pix2struct) + + +## Pix2StructConfig + +[[autodoc]] Pix2StructConfig + - from_text_vision_configs + +## Pix2StructTextConfig + +[[autodoc]] Pix2StructTextConfig + +## Pix2StructVisionConfig + +[[autodoc]] Pix2StructVisionConfig + +## Pix2StructProcessor + +[[autodoc]] Pix2StructProcessor + +## Pix2StructImageProcessor + +[[autodoc]] Pix2StructImageProcessor + - preprocess + +## Pix2StructTextModel + +[[autodoc]] Pix2StructTextModel + - forward + +## Pix2StructVisionModel + +[[autodoc]] Pix2StructVisionModel + - forward + +## Pix2StructForConditionalGeneration + +[[autodoc]] Pix2StructForConditionalGeneration + - forward diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 8bf3655906..d655b21deb 100644 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -403,6 +403,13 @@ _import_structure = { "models.pegasus_x": ["PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP", "PegasusXConfig"], "models.perceiver": ["PERCEIVER_PRETRAINED_CONFIG_ARCHIVE_MAP", "PerceiverConfig", "PerceiverTokenizer"], "models.phobert": ["PhobertTokenizer"], + "models.pix2struct": [ + "PIX2STRUCT_PRETRAINED_CONFIG_ARCHIVE_MAP", + "Pix2StructConfig", + "Pix2StructProcessor", + "Pix2StructTextConfig", + "Pix2StructVisionConfig", + ], "models.plbart": ["PLBART_PRETRAINED_CONFIG_ARCHIVE_MAP", "PLBartConfig"], "models.poolformer": ["POOLFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "PoolFormerConfig"], "models.prophetnet": ["PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP", "ProphetNetConfig", "ProphetNetTokenizer"], @@ -861,6 +868,7 @@ else: _import_structure["models.oneformer"].extend(["OneFormerImageProcessor"]) _import_structure["models.owlvit"].extend(["OwlViTFeatureExtractor", "OwlViTImageProcessor"]) _import_structure["models.perceiver"].extend(["PerceiverFeatureExtractor", "PerceiverImageProcessor"]) + _import_structure["models.pix2struct"].extend(["Pix2StructImageProcessor"]) _import_structure["models.poolformer"].extend(["PoolFormerFeatureExtractor", "PoolFormerImageProcessor"]) _import_structure["models.segformer"].extend(["SegformerFeatureExtractor", "SegformerImageProcessor"]) _import_structure["models.swin2sr"].append("Swin2SRImageProcessor") @@ -2101,6 +2109,15 @@ else: "PerceiverPreTrainedModel", ] ) + _import_structure["models.pix2struct"].extend( + [ + "PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST", + "Pix2StructForConditionalGeneration", + "Pix2StructPreTrainedModel", + "Pix2StructTextModel", + "Pix2StructVisionModel", + ] + ) _import_structure["models.plbart"].extend( [ "PLBART_PRETRAINED_MODEL_ARCHIVE_LIST", @@ -4014,6 +4031,13 @@ if TYPE_CHECKING: from .models.pegasus_x import PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusXConfig from .models.perceiver import PERCEIVER_PRETRAINED_CONFIG_ARCHIVE_MAP, PerceiverConfig, PerceiverTokenizer from .models.phobert import PhobertTokenizer + from .models.pix2struct import ( + PIX2STRUCT_PRETRAINED_CONFIG_ARCHIVE_MAP, + Pix2StructConfig, + Pix2StructProcessor, + Pix2StructTextConfig, + Pix2StructVisionConfig, + ) from .models.plbart import PLBART_PRETRAINED_CONFIG_ARCHIVE_MAP, PLBartConfig from .models.poolformer import POOLFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, PoolFormerConfig from .models.prophetnet import PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP, ProphetNetConfig, ProphetNetTokenizer @@ -4419,6 +4443,7 @@ if TYPE_CHECKING: from .models.oneformer import OneFormerImageProcessor from .models.owlvit import OwlViTFeatureExtractor, OwlViTImageProcessor from .models.perceiver import PerceiverFeatureExtractor, PerceiverImageProcessor + from .models.pix2struct import Pix2StructImageProcessor from .models.poolformer import PoolFormerFeatureExtractor, PoolFormerImageProcessor from .models.segformer import SegformerFeatureExtractor, SegformerImageProcessor from .models.swin2sr import Swin2SRImageProcessor @@ -5435,6 +5460,13 @@ if TYPE_CHECKING: PerceiverModel, PerceiverPreTrainedModel, ) + from .models.pix2struct import ( + PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST, + Pix2StructForConditionalGeneration, + Pix2StructPreTrainedModel, + Pix2StructTextModel, + Pix2StructVisionModel, + ) from .models.plbart import ( PLBART_PRETRAINED_MODEL_ARCHIVE_LIST, PLBartForCausalLM, diff --git a/src/transformers/models/__init__.py b/src/transformers/models/__init__.py index d366eefdcb..d32aa2f6e8 100644 --- a/src/transformers/models/__init__.py +++ b/src/transformers/models/__init__.py @@ -140,6 +140,7 @@ from . import ( pegasus_x, perceiver, phobert, + pix2struct, plbart, poolformer, prophetnet, diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py index 9a8ef79c7c..9a8500d118 100755 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py @@ -141,6 +141,7 @@ CONFIG_MAPPING_NAMES = OrderedDict( ("pegasus", "PegasusConfig"), ("pegasus_x", "PegasusXConfig"), ("perceiver", "PerceiverConfig"), + ("pix2struct", "Pix2StructConfig"), ("plbart", "PLBartConfig"), ("poolformer", "PoolFormerConfig"), ("prophetnet", "ProphetNetConfig"), @@ -315,6 +316,7 @@ CONFIG_ARCHIVE_MAP_MAPPING_NAMES = OrderedDict( ("pegasus", "PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("pegasus_x", "PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("perceiver", "PERCEIVER_PRETRAINED_CONFIG_ARCHIVE_MAP"), + ("pix2struct", "PIX2STRUCT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("plbart", "PLBART_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("poolformer", "POOLFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("prophetnet", "PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP"), @@ -505,6 +507,7 @@ MODEL_NAMES_MAPPING = OrderedDict( ("pegasus_x", "PEGASUS-X"), ("perceiver", "Perceiver"), ("phobert", "PhoBERT"), + ("pix2struct", "Pix2Struct"), ("plbart", "PLBart"), ("poolformer", "PoolFormer"), ("prophetnet", "ProphetNet"), diff --git a/src/transformers/models/auto/image_processing_auto.py b/src/transformers/models/auto/image_processing_auto.py index 95ecd57172..c092dbf16f 100644 --- a/src/transformers/models/auto/image_processing_auto.py +++ b/src/transformers/models/auto/image_processing_auto.py @@ -80,6 +80,7 @@ IMAGE_PROCESSOR_MAPPING_NAMES = OrderedDict( ("oneformer", "OneFormerImageProcessor"), ("owlvit", "OwlViTImageProcessor"), ("perceiver", "PerceiverImageProcessor"), + ("pix2struct", "Pix2StructImageProcessor"), ("poolformer", "PoolFormerImageProcessor"), ("regnet", "ConvNextImageProcessor"), ("resnet", "ConvNextImageProcessor"), diff --git a/src/transformers/models/auto/processing_auto.py b/src/transformers/models/auto/processing_auto.py index fb02003400..9e6edc0ae1 100644 --- a/src/transformers/models/auto/processing_auto.py +++ b/src/transformers/models/auto/processing_auto.py @@ -60,6 +60,7 @@ PROCESSOR_MAPPING_NAMES = OrderedDict( ("mgp-str", "MgpstrProcessor"), ("oneformer", "OneFormerProcessor"), ("owlvit", "OwlViTProcessor"), + ("pix2struct", "Pix2StructProcessor"), ("sew", "Wav2Vec2Processor"), ("sew-d", "Wav2Vec2Processor"), ("speech_to_text", "Speech2TextProcessor"), diff --git a/src/transformers/models/auto/tokenization_auto.py b/src/transformers/models/auto/tokenization_auto.py index a029bfb52c..5afedadca4 100644 --- a/src/transformers/models/auto/tokenization_auto.py +++ b/src/transformers/models/auto/tokenization_auto.py @@ -248,6 +248,7 @@ else: ), ), ("phobert", ("PhobertTokenizer", None)), + ("pix2struct", ("T5Tokenizer", "T5TokenizerFast" if is_tokenizers_available() else None)), ("plbart", ("PLBartTokenizer" if is_sentencepiece_available() else None, None)), ("prophetnet", ("ProphetNetTokenizer", None)), ("qdqbert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)), diff --git a/src/transformers/models/pix2struct/__init__.py b/src/transformers/models/pix2struct/__init__.py new file mode 100644 index 0000000000..8b395b31d8 --- /dev/null +++ b/src/transformers/models/pix2struct/__init__.py @@ -0,0 +1,86 @@ +# Copyright 2023 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. +from typing import TYPE_CHECKING + +from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available + + +_import_structure = { + "configuration_pix2struct": [ + "PIX2STRUCT_PRETRAINED_CONFIG_ARCHIVE_MAP", + "Pix2StructConfig", + "Pix2StructTextConfig", + "Pix2StructVisionConfig", + ], + "processing_pix2struct": ["Pix2StructProcessor"], +} + +try: + if not is_vision_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["image_processing_pix2struct"] = ["Pix2StructImageProcessor"] + + +try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["modeling_pix2struct"] = [ + "PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST", + "Pix2StructPreTrainedModel", + "Pix2StructForConditionalGeneration", + "Pix2StructVisionModel", + "Pix2StructTextModel", + ] + +if TYPE_CHECKING: + from .configuration_pix2struct import ( + PIX2STRUCT_PRETRAINED_CONFIG_ARCHIVE_MAP, + Pix2StructConfig, + Pix2StructTextConfig, + Pix2StructVisionConfig, + ) + from .processing_pix2struct import Pix2StructProcessor + + try: + if not is_vision_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .image_processing_pix2struct import Pix2StructImageProcessor + + try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .modeling_pix2struct import ( + PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST, + Pix2StructForConditionalGeneration, + Pix2StructPreTrainedModel, + Pix2StructTextModel, + Pix2StructVisionModel, + ) + +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/src/transformers/models/pix2struct/configuration_pix2struct.py b/src/transformers/models/pix2struct/configuration_pix2struct.py new file mode 100644 index 0000000000..c8a866966c --- /dev/null +++ b/src/transformers/models/pix2struct/configuration_pix2struct.py @@ -0,0 +1,412 @@ +# coding=utf-8 +# Copyright 2023 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. +""" Pix2Struct model configuration""" + +import copy +import os +from typing import Union + +from ...configuration_utils import PretrainedConfig +from ...utils import logging + + +logger = logging.get_logger(__name__) + +PIX2STRUCT_PRETRAINED_CONFIG_ARCHIVE_MAP = { + "google/pix2struct-textcaps-base": ( + "https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json" + ), +} + + +class Pix2StructTextConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`Pix2StructTextModel`]. It is used to instantiate + a Pix2Struct text model according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the `Pix2StructText` used by the + [base architectures](https://huggingface.co/google/pix2struct-textcaps-base). + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 50244): + Vocabulary size of the `Pix2Struct` text model. Defines the number of different tokens that can be + represented by the `inputs_ids` passed when calling [`Pix2StructModel`]. + hidden_size (`int`, *optional*, defaults to 768): + Dimensionality of the encoder layers and the pooler layer. + d_kv (`int`, *optional*, defaults to 64): + Dimensionality of the key, query, value projections in each attention head. + d_ff (`int`, *optional*, defaults to 2048): + Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. + num_layers (`int`, *optional*, defaults to 12): + Number of hidden layers in the Transformer encoder. + num_heads (`int`, *optional*, defaults to 12): + Number of attention heads for each attention layer in the Transformer encoder. + relative_attention_num_buckets (`int`, *optional*, defaults to 32): + The number of buckets to use for each attention layer. + relative_attention_max_distance (`int`, *optional*, defaults to 128): + The maximum distance of the longer sequences for the bucket separation. + dropout_rate (`float`, *optional*, defaults to 0.1): + The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler. + layer_norm_epsilon (`float`, *optional*, defaults to 1e-6): + The epsilon used by the layer normalization layers. + initializer_factor (`float`, *optional*, defaults to 1.0): + A factor for initializing all weight matrices (should be kept to 1, used internally for initialization + testing). + dense_act_fn (`Union[Callable, str]`, *optional*, defaults to `"gelu_new"`): + The non-linear activation function (function or string). + decoder_start_token_id (`int`, *optional*, defaults to 0): + The id of the `decoder_start_token_id` token. + use_cache (`bool`, *optional*, defaults to `False`): + Whether or not the model should return the last key/values attentions (not used by all models). + pad_token_id (`int`, *optional*, defaults to 0): + The id of the `padding` token. + eos_token_id (`int`, *optional*, defaults to 1): + The id of the `end-of-sequence` token. + + Example: + + ```python + >>> from transformers import Pix2StructTextConfig, Pix2StructTextModel + + >>> # Initializing a Pix2StructTextConfig with Salesforce/pix2struct-vqa-base style configuration + >>> configuration = Pix2StructTextConfig() + + >>> # Initializing a Pix2StructTextModel (with random weights) from the Salesforce/pix2struct-vqa-base style configuration + >>> model = Pix2StructTextModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + model_type = "pix2struct_text_model" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "hidden_size": "hidden_size", + "num_attention_heads": "num_heads", + "num_hidden_layers": "num_layers", + } + + def __init__( + self, + vocab_size=50244, + hidden_size=768, + d_kv=64, + d_ff=2048, + num_layers=12, + num_heads=12, + relative_attention_num_buckets=32, + relative_attention_max_distance=128, + dropout_rate=0.1, + layer_norm_epsilon=1e-6, + initializer_factor=1.0, + dense_act_fn="gelu_new", + decoder_start_token_id=0, + use_cache=False, + pad_token_id=0, + eos_token_id=1, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.d_kv = d_kv + self.d_ff = d_ff + self.num_layers = num_layers + self.num_heads = num_heads + self.relative_attention_num_buckets = relative_attention_num_buckets + self.relative_attention_max_distance = relative_attention_max_distance + self.dropout_rate = dropout_rate + self.layer_norm_epsilon = layer_norm_epsilon + self.initializer_factor = initializer_factor + self.use_cache = use_cache + + self.eos_token_id = eos_token_id + self.decoder_start_token_id = decoder_start_token_id + + # for backwards compatibility + self.dense_act_fn = dense_act_fn + + super().__init__( + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + decoder_start_token_id=decoder_start_token_id, + **kwargs, + ) + + @classmethod + def from_pretrained( + cls, pretrainehidden_size_name_or_path: Union[str, os.PathLike], **kwargs + ) -> "PretrainedConfig": + config_dict, kwargs = cls.get_config_dict(pretrainehidden_size_name_or_path, **kwargs) + + # get the text config dict if we are loading from Pix2StructConfig + if config_dict.get("model_type") == "pix2struct": + config_dict = config_dict["text_config"] + + if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type: + logger.warning( + f"You are using a model of type {config_dict['model_type']} to instantiate a model of type " + f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." + ) + + return cls.from_dict(config_dict, **kwargs) + + +class Pix2StructVisionConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`Pix2StructVisionModel`]. It is used to + instantiate a PIX2STRUCT vision model according to the specified arguments, defining the model architecture. + Instantiating a configuration defaults will yield a similar configuration to that of the Pix2Struct-base + [Salesforce/pix2struct-vqa-base](https://huggingface.co/Salesforce/pix2struct-vqa-base) 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. + patch_embed_hidden_size (`int`, *optional*, defaults to 768): + Dimensionality of the input patch_embedding layer in the Transformer encoder. + d_ff (`int`, *optional*, defaults to 2048): + Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. + d_kv (`int`, *optional*, defaults to 64): + Dimensionality of the key, query, value projections per attention head. + projection_dim (`int`, *optional*, defaults to 768): + Dimensionality of the projection layer in the Transformer 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 of the input images. + patch_size (`int`, *optional*, defaults to 16): + The size (resolution) of each patch. + dense_act_fn (`str` or `function`, *optional*, defaults to `"gelu_new"`): + The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, + `"relu"`, `"selu"` and `"gelu_new"` ``"gelu"` are supported. + layer_norm_eps (`float`, *optional*, defaults to 1e-6): + The epsilon used by the layer normalization layers. + dropout_rate (`float`, *optional*, defaults to 0.0): + The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler. + 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. + initializer_factor (`float``, *optional*, defaults to 1): + A factor for initializing all weight matrices (should be kept to 1, used internally for initialization + testing). + seq_len (`int`, *optional*, defaults to 4096): + Maximum sequence length (here number of patches) supported by the model. + layer_norm_bias (`bool`, *optional*, defaults to `False`): + Whether or not to add a bias to the layer normalization layers. + relative_attention_num_buckets (`int`, *optional*, defaults to 32): + The number of buckets to use for each attention layer. + relative_attention_max_distance (`int`, *optional*, defaults to 128): + The maximum distance (in tokens) to use for each attention layer. + + Example: + + ```python + >>> from transformers import Pix2StructVisionConfig, Pix2StructVisionModel + + >>> # Initializing a Pix2StructVisionConfig with Salesforce/pix2struct-vqa-base style configuration + >>> configuration = Pix2StructVisionConfig() + + >>> # Initializing a Pix2StructVisionModel (with random weights) from the Salesforce/pix2struct-vqa-base style configuration + >>> model = Pix2StructVisionModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "pix2struct_vision_model" + + def __init__( + self, + hidden_size=768, + patch_embed_hidden_size=768, + d_ff=2048, + d_kv=64, + projection_dim=768, + num_hidden_layers=12, + num_attention_heads=12, + num_channels=3, + patch_size=16, + dense_act_fn="gelu_new", + layer_norm_eps=1e-6, + dropout_rate=0.0, + attention_dropout=0.0, + initializer_range=1e-10, + initializer_factor=1.0, + seq_len=4096, + layer_norm_bias=False, + relative_attention_num_buckets=32, + relative_attention_max_distance=128, + **kwargs, + ): + super().__init__(**kwargs) + + self.hidden_size = hidden_size + self.patch_embed_hidden_size = patch_embed_hidden_size + self.d_ff = d_ff + self.projection_dim = projection_dim + self.dropout_rate = dropout_rate + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_channels = num_channels + self.patch_size = patch_size + self.initializer_range = initializer_range + self.initializer_factor = initializer_factor + self.attention_dropout = attention_dropout + self.layer_norm_eps = layer_norm_eps + self.dense_act_fn = dense_act_fn + self.seq_len = seq_len + self.layer_norm_bias = layer_norm_bias + self.relative_attention_num_buckets = relative_attention_num_buckets + self.relative_attention_max_distance = relative_attention_max_distance + self.d_kv = d_kv + + @classmethod + def from_pretrained( + cls, pretrainehidden_size_name_or_path: Union[str, os.PathLike], **kwargs + ) -> "PretrainedConfig": + config_dict, kwargs = cls.get_config_dict(pretrainehidden_size_name_or_path, **kwargs) + + # get the vision config dict if we are loading from Pix2StructConfig + if config_dict.get("model_type") == "pix2struct": + config_dict = config_dict["vision_config"] + + if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type: + logger.warning( + f"You are using a model of type {config_dict['model_type']} to instantiate a model of type " + f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." + ) + + return cls.from_dict(config_dict, **kwargs) + + +class Pix2StructConfig(PretrainedConfig): + r""" + [`Pix2StructConfig`] is the configuration class to store the configuration of a [`Pix2StructModel`]. It is used to + instantiate a PIX2STRUCT model according to the specified arguments, defining the text model and vision model + configs. Instantiating a configuration with the defaults will yield a similar configuration to that of the + PIX2STRUCT-base [Salesforce/pix2struct-vqa-base](https://huggingface.co/Salesforce/pix2struct-vqa-base) + architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + text_config (`dict`, *optional*): + Dictionary of configuration options used to initialize [`Pix2StructTextConfig`]. + vision_config (`dict`, *optional*): + Dictionary of configuration options used to initialize [`Pix2StructVisionConfig`]. + initializer_factor (`float`, *optional*, defaults to 1.0): + Factor to multiply the initialization range with. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + is_vqa (`bool`, *optional*, defaults to `False`): + Whether the model has been fine-tuned for VQA or not. + kwargs (*optional*): + Dictionary of keyword arguments. + + Example: + + ```python + >>> from transformers import Pix2StructConfig, Pix2StructModel + + >>> # Initializing a Pix2StructConfig with Salesforce/pix2struct-vqa-base style configuration + >>> configuration = Pix2StructConfig() + + >>> # Initializing a Pix2StructPModel (with random weights) from the Salesforce/pix2struct-vqa-base style configuration + >>> model = Pix2StructModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + + >>> # We can also initialize a Pix2StructConfig from a Pix2StructTextConfig and a Pix2StructVisionConfig + + >>> # Initializing a PIX2STRUCTText and PIX2STRUCTVision configuration + >>> config_text = Pix2StructTextConfig() + >>> config_vision = Pix2StructVisionConfig() + + >>> config = Pix2StructConfig.from_text_vision_configs(config_text, config_vision) + ```""" + + model_type = "pix2struct" + is_composition = True + + def __init__( + self, + text_config=None, + vision_config=None, + initializer_factor=1.0, + initializer_range=0.02, + is_vqa=False, + **kwargs, + ): + super().__init__(**kwargs) + + if text_config is None: + text_config = {} + logger.info("text_config is None. Initializing the Pix2StructTextConfig with default values.") + + if vision_config is None: + vision_config = {} + logger.info("vision_config is None. Initializing the Pix2StructVisionConfig with default values.") + + self.text_config = Pix2StructTextConfig(**text_config) + self.vision_config = Pix2StructVisionConfig(**vision_config) + + self.text_config.encoder_hidden_size = self.vision_config.hidden_size + self.decoder_start_token_id = self.text_config.decoder_start_token_id + self.pad_token_id = self.text_config.pad_token_id + + self.initializer_factor = initializer_factor + self.initializer_range = initializer_range + + self.text_config.initializer_range = self.initializer_range + self.vision_config.initializer_range = self.initializer_range + + self.is_vqa = is_vqa + + @classmethod + def from_text_vision_configs( + cls, text_config: Pix2StructTextConfig, vision_config: Pix2StructVisionConfig, **kwargs + ): + r""" + Instantiate a [`Pix2StructConfig`] (or a derived class) from pix2struct text model configuration and pix2struct + vision model configuration. + + Returns: + [`Pix2StructConfig`]: An instance of a configuration object + """ + + return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs) + + def to_dict(self): + """ + Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. + + Returns: + `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance, + """ + output = copy.deepcopy(self.__dict__) + output["text_config"] = self.text_config.to_dict() + output["vision_config"] = self.vision_config.to_dict() + output["model_type"] = self.__class__.model_type + return output diff --git a/src/transformers/models/pix2struct/convert_pix2struct_original_pytorch_to_hf.py b/src/transformers/models/pix2struct/convert_pix2struct_original_pytorch_to_hf.py new file mode 100644 index 0000000000..457c223669 --- /dev/null +++ b/src/transformers/models/pix2struct/convert_pix2struct_original_pytorch_to_hf.py @@ -0,0 +1,155 @@ +# coding=utf-8 +# Copyright 2023 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 os +import re + +import torch +from flax.traverse_util import flatten_dict +from t5x import checkpoints + +from transformers import ( + AutoTokenizer, + Pix2StructConfig, + Pix2StructForConditionalGeneration, + Pix2StructImageProcessor, + Pix2StructProcessor, + Pix2StructTextConfig, + Pix2StructVisionConfig, +) + + +def get_flax_param(t5x_checkpoint_path): + flax_params = checkpoints.load_t5x_checkpoint(t5x_checkpoint_path) + flax_params = flatten_dict(flax_params) + return flax_params + + +def rename_and_convert_flax_params(flax_dict): + converted_dict = {} + + CONVERSION_MAPPING = { + "token_embedder": "embeddings", + "encoder_norm": "layernorm", + "kernel": "weight", + ".out": ".output", + "scale": "weight", + "embedders_0.pos_embedding": "row_embedder.weight", + "embedders_1.pos_embedding": "column_embedder.weight", + } + + DECODER_CONVERSION_MAPPING = { + "query": "attention.query", + "key": "attention.key", + "value": "attention.value", + "output.dense": "output", + "encoder_decoder_attention.o": "encoder_decoder_attention.attention.o", + "pre_self_attention_layer_norm": "self_attention.layer_norm", + "pre_cross_attention_layer_norm": "encoder_decoder_attention.layer_norm", + "mlp.": "mlp.DenseReluDense.", + "pre_mlp_layer_norm": "mlp.layer_norm", + "self_attention.o": "self_attention.attention.o", + "decoder.embeddings.embedding": "decoder.embed_tokens.weight", + "decoder.relpos_bias.rel_embedding": "decoder.layer.0.self_attention.attention.relative_attention_bias.weight", + "decoder.decoder_norm.weight": "decoder.final_layer_norm.weight", + "decoder.logits_dense.weight": "decoder.lm_head.weight", + } + + for key in flax_dict.keys(): + if "target" in key: + # remove the first prefix from the key + new_key = ".".join(key[1:]) + + # rename the key + for old, new in CONVERSION_MAPPING.items(): + new_key = new_key.replace(old, new) + + if "decoder" in new_key: + for old, new in DECODER_CONVERSION_MAPPING.items(): + new_key = new_key.replace(old, new) + + if "layers" in new_key and "decoder" not in new_key: + # use regex to replace the layer number + new_key = re.sub(r"layers_(\d+)", r"layer.\1", new_key) + new_key = new_key.replace("encoder", "encoder.encoder") + + elif "layers" in new_key and "decoder" in new_key: + # use regex to replace the layer number + new_key = re.sub(r"layers_(\d+)", r"layer.\1", new_key) + + converted_dict[new_key] = flax_dict[key] + + converted_torch_dict = {} + # convert converted_dict into torch format + for key in converted_dict.keys(): + if ("embed_tokens" not in key) and ("embedder" not in key): + converted_torch_dict[key] = torch.from_numpy(converted_dict[key].T) + else: + converted_torch_dict[key] = torch.from_numpy(converted_dict[key]) + + return converted_torch_dict + + +def convert_pix2struct_original_pytorch_checkpoint_to_hf( + t5x_checkpoint_path, pytorch_dump_folder_path, use_large=False, is_vqa=False +): + flax_params = get_flax_param(t5x_checkpoint_path) + + if not use_large: + encoder_config = Pix2StructVisionConfig() + decoder_config = Pix2StructTextConfig() + else: + encoder_config = Pix2StructVisionConfig( + hidden_size=1536, d_ff=3968, num_attention_heads=24, num_hidden_layers=18 + ) + decoder_config = Pix2StructTextConfig(hidden_size=1536, d_ff=3968, num_heads=24, num_layers=18) + config = Pix2StructConfig( + vision_config=encoder_config.to_dict(), text_config=decoder_config.to_dict(), is_vqa=is_vqa + ) + + model = Pix2StructForConditionalGeneration(config) + + torch_params = rename_and_convert_flax_params(flax_params) + model.load_state_dict(torch_params) + + tok = AutoTokenizer.from_pretrained("ybelkada/test-pix2struct-tokenizer") + image_processor = Pix2StructImageProcessor() + processor = Pix2StructProcessor(image_processor=image_processor, tokenizer=tok) + + if use_large: + processor.image_processor.max_patches = 4096 + + processor.image_processor.is_vqa = True + + # mkdir if needed + os.makedirs(pytorch_dump_folder_path, exist_ok=True) + + model.save_pretrained(pytorch_dump_folder_path) + processor.save_pretrained(pytorch_dump_folder_path) + + print("Model saved in {}".format(pytorch_dump_folder_path)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--t5x_checkpoint_path", default=None, type=str, help="Path to the original T5x checkpoint.") + parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.") + parser.add_argument("--use_large", action="store_true", help="Use large model.") + parser.add_argument("--is_vqa", action="store_true", help="Use large model.") + args = parser.parse_args() + + convert_pix2struct_original_pytorch_checkpoint_to_hf( + args.t5x_checkpoint_path, args.pytorch_dump_folder_path, args.use_large + ) diff --git a/src/transformers/models/pix2struct/image_processing_pix2struct.py b/src/transformers/models/pix2struct/image_processing_pix2struct.py new file mode 100644 index 0000000000..b85d123841 --- /dev/null +++ b/src/transformers/models/pix2struct/image_processing_pix2struct.py @@ -0,0 +1,418 @@ +# coding=utf-8 +# Copyright 2023 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. +"""Image processor class for Pix2Struct.""" +import io +import math +from typing import Dict, Optional, Union + +import numpy as np +from huggingface_hub import hf_hub_download + +from ...image_processing_utils import BaseImageProcessor, BatchFeature +from ...image_transforms import convert_to_rgb, normalize, to_channel_dimension_format, to_pil_image +from ...image_utils import ( + ChannelDimension, + ImageInput, + get_image_size, + infer_channel_dimension_format, + make_list_of_images, + to_numpy_array, + valid_images, +) +from ...utils import TensorType, is_torch_available, is_vision_available, logging +from ...utils.import_utils import requires_backends + + +if is_vision_available(): + import textwrap + + from PIL import Image, ImageDraw, ImageFont + +if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) +DEFAULT_FONT_PATH = "ybelkada/fonts" + + +# adapted from: https://discuss.pytorch.org/t/tf-image-extract-patches-in-pytorch/171409/2 +def torch_extract_patches(image_tensor, patch_height, patch_width): + """ + Utiliy function to extract patches from a given image tensor. Returns a tensor of shape (1, `patch_height`, + `patch_width`, `num_channels`x `patch_height` x `patch_width`) + + Args: + image_tensor (torch.Tensor): + The image tensor to extract patches from. + patch_height (int): + The height of the patches to extract. + patch_width (int): + The width of the patches to extract. + """ + requires_backends(torch_extract_patches, ["torch"]) + + image_tensor = image_tensor.unsqueeze(0) + patches = torch.nn.functional.unfold(image_tensor, (patch_height, patch_width), stride=(patch_height, patch_width)) + patches = patches.reshape(image_tensor.size(0), image_tensor.size(1), patch_height, patch_width, -1) + patches = patches.permute(0, 4, 2, 3, 1).reshape( + image_tensor.size(2) // patch_height, + image_tensor.size(3) // patch_width, + image_tensor.size(1) * patch_height * patch_width, + ) + return patches.unsqueeze(0) + + +# Adapted from https://github.com/google-research/pix2struct/blob/0e1779af0f4db4b652c1d92b3bbd2550a7399123/pix2struct/preprocessing/preprocessing_utils.py#L106 +def render_text( + text: str, + text_size: int = 36, + text_color: str = "black", + background_color: str = "white", + left_padding: int = 5, + right_padding: int = 5, + top_padding: int = 5, + bottom_padding: int = 5, + font_bytes: Optional[bytes] = None, + font_path: Optional[str] = None, +) -> Image.Image: + """ + Render text. This script is entirely adapted from the original script that can be found here: + https://github.com/google-research/pix2struct/blob/main/pix2struct/preprocessing/preprocessing_utils.py + + Args: + text (`str`, *optional*, defaults to ): + Text to render. + text_size (`int`, *optional*, defaults to 36): + Size of the text. + text_color (`str`, *optional*, defaults to `"black"`): + Color of the text. + background_color (`str`, *optional*, defaults to `"white"`): + Color of the background. + left_padding (`int`, *optional*, defaults to 5): + Padding on the left. + right_padding (`int`, *optional*, defaults to 5): + Padding on the right. + top_padding (`int`, *optional*, defaults to 5): + Padding on the top. + bottom_padding (`int`, *optional*, defaults to 5): + Padding on the bottom. + font_bytes (`bytes`, *optional*): + Bytes of the font to use. If `None`, the default font will be used. + font_path (`str`, *optional*): + Path to the font to use. If `None`, the default font will be used. + """ + requires_backends(render_text, "vision") + # Add new lines so that each line is no more than 80 characters. + + wrapper = textwrap.TextWrapper(width=80) + lines = wrapper.wrap(text=text) + wrapped_text = "\n".join(lines) + + if font_bytes is not None and font_path is None: + font = io.BytesIO(font_bytes) + elif font_path is not None: + font = font_path + else: + font = hf_hub_download(DEFAULT_FONT_PATH, "Arial.TTF") + font = ImageFont.truetype(font, encoding="UTF-8", size=text_size) + + # Use a temporary canvas to determine the width and height in pixels when + # rendering the text. + temp_draw = ImageDraw.Draw(Image.new("RGB", (1, 1), background_color)) + _, _, text_width, text_height = temp_draw.textbbox((0, 0), wrapped_text, font) + + # Create the actual image with a bit of padding around the text. + image_width = text_width + left_padding + right_padding + image_height = text_height + top_padding + bottom_padding + image = Image.new("RGB", (image_width, image_height), background_color) + draw = ImageDraw.Draw(image) + draw.text(xy=(left_padding, top_padding), text=wrapped_text, fill=text_color, font=font) + return image + + +# Adapted from https://github.com/google-research/pix2struct/blob/0e1779af0f4db4b652c1d92b3bbd2550a7399123/pix2struct/preprocessing/preprocessing_utils.py#L87 +def render_header(image: np.ndarray, header: str, **kwargs): + """ + Renders the input text as a header on the input image. + + Args: + image (`np.ndarray`): + The image to render the header on. + header (`str`): + The header text. + data_format (`Union[ChannelDimension, str]`, *optional*): + The data format of the image. Can be either "ChannelDimension.channels_first" or + "ChannelDimension.channels_last". + + Returns: + `np.ndarray`: The image with the header rendered. + """ + requires_backends(render_header, "vision") + + # Convert to PIL image if necessary + image = to_pil_image(image) + + header_image = render_text(header, **kwargs) + new_width = max(header_image.width, image.width) + + new_height = int(image.height * (new_width / image.width)) + new_header_height = int(header_image.height * (new_width / header_image.width)) + + new_image = Image.new("RGB", (new_width, new_height + new_header_height), "white") + new_image.paste(header_image.resize((new_width, new_header_height)), (0, 0)) + new_image.paste(image.resize((new_width, new_height)), (0, new_header_height)) + + # Convert back to the original framework if necessary + new_image = to_numpy_array(new_image) + + if infer_channel_dimension_format(new_image) == ChannelDimension.LAST: + new_image = to_channel_dimension_format(new_image, ChannelDimension.LAST) + + return new_image + + +class Pix2StructImageProcessor(BaseImageProcessor): + r""" + Constructs a Pix2Struct image processor. + + Args: + do_convert_rgb (`bool`, *optional*, defaults to `True`): + Whether to convert the image to RGB. + do_normalize (`bool`, *optional*, defaults to `True`): + Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess` + method. According to Pix2Struct paper and code, the image is normalized with its own mean and standard + deviation. + patch_size (`Dict[str, int]`, *optional*, defaults to `{"height": 16, "width": 16}`): + The patch size to use for the image. According to Pix2Struct paper and code, the patch size is 16x16. + max_patches (`int`, *optional*, defaults to 2048): + The maximum number of patches to extract from the image as per the [Pix2Struct + paper](https://arxiv.org/pdf/2210.03347.pdf). + is_vqa (`bool`, *optional*, defaults to `False`): + Whether or not the image processor is for the VQA task. If `True` and `header_text` is passed in, text is + rendered onto the input images. + """ + + model_input_names = ["flattened_patches"] + + def __init__( + self, + do_convert_rgb: bool = True, + do_normalize: bool = True, + patch_size: Dict[str, int] = None, + max_patches: int = 2048, + is_vqa: bool = False, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.patch_size = patch_size if patch_size is not None else {"height": 16, "width": 16} + self.do_normalize = do_normalize + self.do_convert_rgb = do_convert_rgb + self.max_patches = max_patches + self.is_vqa = is_vqa + + def extract_flattened_patches(self, image: np.ndarray, max_patches: int, patch_size: dict, **kwargs) -> np.ndarray: + """ + Extract flattened patches from an image. + + Args: + image (`np.ndarray`): + Image to extract flattened patches from. + max_patches (`int`): + Maximum number of patches to extract. + patch_size (`dict`): + Dictionary containing the patch height and width. + + Returns: + result (`np.ndarray`): + A sequence of `max_patches` flattened patches. + """ + requires_backends(self.extract_flattened_patches, "torch") + + # convert to torch + image = to_channel_dimension_format(image, ChannelDimension.FIRST) + image = torch.from_numpy(image) + + patch_height, patch_width = patch_size["height"], patch_size["width"] + image_height, image_width = get_image_size(image) + + # maximize scale s.t. + scale = math.sqrt(max_patches * (patch_height / image_height) * (patch_width / image_width)) + num_feasible_rows = max(min(math.floor(scale * image_height / patch_height), max_patches), 1) + num_feasible_cols = max(min(math.floor(scale * image_width / patch_width), max_patches), 1) + resized_height = max(num_feasible_rows * patch_height, 1) + resized_width = max(num_feasible_cols * patch_width, 1) + + image = torch.nn.functional.interpolate( + image.unsqueeze(0), + size=(resized_height, resized_width), + mode="bilinear", + align_corners=False, + antialias=True, + ).squeeze(0) + + # [1, rows, columns, patch_height * patch_width * image_channels] + patches = torch_extract_patches(image, patch_height, patch_width) + + patches_shape = patches.shape + rows = patches_shape[1] + columns = patches_shape[2] + depth = patches_shape[3] + + # [rows * columns, patch_height * patch_width * image_channels] + patches = patches.reshape([rows * columns, depth]) + + # [rows * columns, 1] + row_ids = torch.arange(rows).reshape([rows, 1]).repeat(1, columns).reshape([rows * columns, 1]) + col_ids = torch.arange(columns).reshape([1, columns]).repeat(rows, 1).reshape([rows * columns, 1]) + + # Offset by 1 so the ids do not contain zeros, which represent padding. + row_ids += 1 + col_ids += 1 + + # Prepare additional patch features. + # [rows * columns, 1] + row_ids = row_ids.to(torch.float32) + col_ids = col_ids.to(torch.float32) + + # [rows * columns, 2 + patch_height * patch_width * image_channels] + result = torch.cat([row_ids, col_ids, patches], -1) + + # [max_patches, 2 + patch_height * patch_width * image_channels] + result = torch.nn.functional.pad(result, [0, 0, 0, max_patches - (rows * columns)]).float() + + result = to_numpy_array(result) + + return result + + def normalize( + self, image: np.ndarray, data_format: Optional[Union[str, ChannelDimension]] = None, **kwargs + ) -> np.ndarray: + """ + Normalize an image. image = (image - image_mean) / image_std. + + The image std is to mimic the tensorflow implementation of the `per_image_standardization`: + https://www.tensorflow.org/api_docs/python/tf/image/per_image_standardization + + Args: + image (`np.ndarray`): + Image to normalize. + """ + if image.dtype == np.uint8: + image = image.astype(np.float32) + + # take mean across the whole `image` + mean = np.mean(image) + std = np.std(image) + adjusted_stddev = max(std, 1.0 / math.sqrt(np.prod(image.shape))) + + return normalize(image, mean=mean, std=adjusted_stddev, **kwargs) + + def preprocess( + self, + images: ImageInput, + header_text: Optional[str] = None, + do_convert_rgb: bool = None, + do_normalize: Optional[bool] = None, + max_patches: Optional[int] = None, + patch_size: Optional[Dict[str, int]] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + data_format: ChannelDimension = ChannelDimension.FIRST, + **kwargs, + ) -> ImageInput: + """ + Preprocess an image or batch of images. The processor first computes the maximum possible number of + aspect-ratio preserving patches of size `patch_size` that can be extracted from the image. It then pads the + image with zeros to make the image respect the constraint of `max_patches`. Before extracting the patches the + images are standardized following the tensorflow implementation of `per_image_standardization` + (https://www.tensorflow.org/api_docs/python/tf/image/per_image_standardization). + + + Args: + images (`ImageInput`): + Image to preprocess. + header_text (`Union[List[str], str]`, *optional*): + Text to render as a header. Only has an effect if `image_processor.is_vqa` is `True`. + do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`): + Whether to convert the image to RGB. + do_normalize (`bool`, *optional*, defaults to `self.do_normalize`): + Whether to normalize the image. + max_patches (`int`, *optional*, defaults to `self.max_patches`): + Maximum number of patches to extract. + patch_size (`dict`, *optional*, defaults to `self.patch_size`): + Dictionary containing the patch height and width. + 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`. + """ + do_normalize = do_normalize if do_normalize is not None else self.do_normalize + do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb + patch_size = patch_size if patch_size is not None else self.patch_size + max_patches = max_patches if max_patches is not None else self.max_patches + is_vqa = self.is_vqa + + if kwargs.get("data_format", None) is not None: + raise ValueError("data_format is not an accepted input as the outputs are ") + + images = make_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." + ) + + # 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 is_vqa: + if header_text is None: + raise ValueError("A header text must be provided for VQA models.") + font_bytes = kwargs.pop("font_bytes", None) + font_path = kwargs.pop("font_path", None) + + if isinstance(header_text, str): + header_text = [header_text] * len(images) + + images = [ + render_header(image, header_text[i], font_bytes=font_bytes, font_path=font_path) + for i, image in enumerate(images) + ] + + if do_normalize: + images = [self.normalize(image=image) for image in images] + + # convert to torch tensor and permute + images = [ + self.extract_flattened_patches(image=image, max_patches=max_patches, patch_size=patch_size) + for image in images + ] + + # create attention mask in numpy + attention_masks = [(image.sum(axis=-1) != 0).astype(np.float32) for image in images] + + encoded_outputs = BatchFeature( + data={"flattened_patches": images, "attention_mask": attention_masks}, tensor_type=return_tensors + ) + + return encoded_outputs diff --git a/src/transformers/models/pix2struct/modeling_pix2struct.py b/src/transformers/models/pix2struct/modeling_pix2struct.py new file mode 100644 index 0000000000..1dcd22e0f4 --- /dev/null +++ b/src/transformers/models/pix2struct/modeling_pix2struct.py @@ -0,0 +1,1823 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. & Google 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. +""" Pix2Struct modeling file""" + +import copy +import math +from typing import Dict, List, Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import nn +from torch.utils.checkpoint import checkpoint + +from ...activations import ACT2FN +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPooling, + CausalLMOutputWithCrossAttentions, + Seq2SeqLMOutput, + Seq2SeqModelOutput, +) +from ...modeling_utils import PreTrainedModel +from ...pytorch_utils import ALL_LAYERNORM_LAYERS +from ...utils import ( + DUMMY_INPUTS, + DUMMY_MASK, + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_torch_fx_proxy, + logging, + replace_return_docstrings, +) +from .configuration_pix2struct import Pix2StructConfig, Pix2StructTextConfig, Pix2StructVisionConfig + + +logger = logging.get_logger(__name__) + +# General docstring +_CONFIG_FOR_DOC = "Pix2StructConfig" + + +PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "google/pix2struct-textcaps-base", + "google/pix2struct-textcaps-large", + "google/pix2struct-base", + "google/pix2struct-large", + "google/pix2struct-ai2d-base", + "google/pix2struct-ai2d-large", + "google/pix2struct-widget-captioning-base", + "google/pix2struct-widget-captioning-large", + "google/pix2struct-screen2words-base", + "google/pix2struct-screen2words-large", + "google/pix2struct-docvqa-base", + "google/pix2struct-docvqa-large", + "google/pix2struct-ocrvqa-base", + "google/pix2struct-ocrvqa-large", + "google/pix2struct-chartqa-base", + "google/pix2struct-inforgraphics-vqa-base", + "google/pix2struct-inforgraphics-vqa-large", + # See all Pix2StructVision models at https://huggingface.co/models?filter=pix2struct +] + + +# Adapted from transformers.models.t5.modeling_t5.T5LayerNorm with T5->Pix2Struct +class Pix2StructLayerNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Construct a layernorm module in the T5 style. No bias and no subtraction of mean. + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + # T5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean + # Square Layer Normalization https://arxiv.org/abs/1910.07467 thus varience is calculated + # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for + # half-precision inputs is done in fp32 + + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + + # convert into half-precision if necessary + if self.weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(self.weight.dtype) + + return self.weight * hidden_states + + +try: + from apex.normalization import FusedRMSNorm + + Pix2StructLayerNorm = FusedRMSNorm # noqa + + logger.info("Discovered apex.normalization.FusedRMSNorm - will use it instead of Pix2StructLayerNorm") +except ImportError: + # using the normal Pix2StructLayerNorm + pass +except Exception: + logger.warning("Discovered apex but it failed to load, falling back to Pix2StructLayerNorm") + pass + +ALL_LAYERNORM_LAYERS.append(Pix2StructLayerNorm) + + +class Pix2StructVisionEmbeddings(nn.Module): + r""" + Construct the embeddings from patch. In `Pix2Struct` the input is different from classic Vision-transformer models. + Here the input is a sequence of `seq_len` flattened patches that also combines padding patches (tokens). Each patch + is represented by a vector of `hidden_size` values. + """ + + def __init__(self, config: Pix2StructConfig) -> None: + super().__init__() + self.patch_projection = nn.Linear(config.patch_embed_hidden_size, config.hidden_size) + + self.row_embedder = nn.Embedding(config.seq_len, config.hidden_size) + self.column_embedder = nn.Embedding(config.seq_len, config.hidden_size) + + self.dropout = nn.Dropout(config.dropout_rate) + + def forward(self, flattened_patches: torch.Tensor) -> torch.Tensor: + # the row and column indices are stored in the first and second position of the flattened_patches + # flattened_patches: `batch_size`, `seq_len`, `hidden_size` + 2 + row_indices = flattened_patches[:, :, 0].long() + col_indices = flattened_patches[:, :, 1].long() + + flattened_patches = flattened_patches[:, :, 2:] + + embeddings = self.patch_projection(flattened_patches) + row_embeddings = self.row_embedder(row_indices) + col_embeddings = self.column_embedder(col_indices) + + # sum all embeddings together + embeddings = embeddings + row_embeddings + col_embeddings + + embeddings = self.dropout(embeddings) + + return embeddings + + +class Pix2StructVisionAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_size = config.hidden_size + self.key_value_proj_dim = config.d_kv + self.n_heads = config.num_attention_heads + self.dropout = config.attention_dropout + self.inner_dim = self.n_heads * self.key_value_proj_dim + + # Mesh TensorFlow initialization to avoid scaling before softmax + self.query = nn.Linear(self.hidden_size, self.inner_dim, bias=False) + self.key = nn.Linear(self.hidden_size, self.inner_dim, bias=False) + self.value = nn.Linear(self.hidden_size, self.inner_dim, bias=False) + self.output = nn.Linear(self.inner_dim, self.hidden_size, bias=False) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + attention_mask=None, + position_bias=None, + layer_head_mask=None, + output_attentions=False, + ): + """ + Self-attention block + """ + # Input is (batch_size, seq_length, dim) + # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length) + # past_key_value[0] is (batch_size, n_heads, q_len - 1, dim_per_head) + batch_size, seq_length = hidden_states.shape[:2] + + def to_projection_shape(states): + """projection""" + return states.contiguous().view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) + + # get query states + # (batch_size, n_heads, seq_length, dim_per_head) + query_states = to_projection_shape(self.query(hidden_states)) + + # get key/value states + key_states = to_projection_shape(self.key(hidden_states)) + value_states = to_projection_shape(self.value(hidden_states)) + + # compute scores + # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 + scores = torch.matmul(query_states, key_states.transpose(3, 2)) + + if position_bias is None: + position_bias = torch.zeros( + (1, self.n_heads, seq_length, seq_length), device=scores.device, dtype=scores.dtype + ) + if self.gradient_checkpointing and self.training: + position_bias.requires_grad = True + + if attention_mask is None: + attention_mask = torch.ones((batch_size, seq_length), device=scores.device, dtype=scores.dtype) + + if attention_mask.dim() == 2: + position_bias = position_bias + attention_mask[:, None, :, None].to(position_bias.device) + else: + # (batch_size, n_heads, seq_length, key_length) + position_bias = position_bias + attention_mask.to(position_bias.device) + position_bias = 1 - position_bias + + position_bias_masked = position_bias.masked_fill(position_bias == 1, torch.finfo(scores.dtype).min) + scores += position_bias_masked + scores = torch.max(scores, torch.tensor(torch.finfo(scores.dtype).min)) + + # (batch_size, n_heads, seq_length, key_length) + attn_weights = nn.functional.softmax(scores, dim=-1, dtype=torch.float32).type_as(scores) + + # (batch_size, n_heads, seq_length, key_length) + attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + + # Mask heads if we want to + if layer_head_mask is not None: + attn_weights = attn_weights * layer_head_mask + + attn_output = torch.matmul(attn_weights, value_states) + + # (batch_size, seq_length, dim) + attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim) + + attn_output = self.output(attn_output) + + outputs = (attn_output,) + (position_bias,) + + if output_attentions: + outputs = outputs + (attn_weights,) + return outputs + + +# Copied from transformers.models.t5.modeling_t5.T5DenseGatedActDense with T5DenseGatedActDense->Pix2StructVisionMlp,T5Config->Pix2StructVisionConfig,config.d_model->config.hidden_size,dropout_rate->dropout_rate +class Pix2StructVisionMlp(nn.Module): + def __init__(self, config: Pix2StructVisionConfig): + super().__init__() + self.wi_0 = nn.Linear(config.hidden_size, config.d_ff, bias=False) + self.wi_1 = nn.Linear(config.hidden_size, config.d_ff, bias=False) + self.wo = nn.Linear(config.d_ff, config.hidden_size, bias=False) + self.dropout = nn.Dropout(config.dropout_rate) + self.act = ACT2FN[config.dense_act_fn] + + def forward(self, hidden_states): + hidden_gelu = self.act(self.wi_0(hidden_states)) + hidden_linear = self.wi_1(hidden_states) + hidden_states = hidden_gelu * hidden_linear + hidden_states = self.dropout(hidden_states) + + # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32. + # See https://github.com/huggingface/transformers/issues/20287 + # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None`` + if ( + isinstance(self.wo.weight, torch.Tensor) + and hidden_states.dtype != self.wo.weight.dtype + and self.wo.weight.dtype != torch.int8 + ): + hidden_states = hidden_states.to(self.wo.weight.dtype) + + hidden_states = self.wo(hidden_states) + return hidden_states + + +class Pix2StructVisionLayer(nn.Module): + def __init__(self, config: Pix2StructConfig) -> None: + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = Pix2StructVisionAttention(config) + self.mlp = Pix2StructVisionMlp(config) + self.pre_mlp_layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.pre_attention_layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + residual = hidden_states + + # in Pix2StructVision, layernorm is applied before self-attention + hidden_states = self.pre_attention_layer_norm(hidden_states) + + self_attention_outputs = self.attention( + hidden_states, + attention_mask=attention_mask, + layer_head_mask=head_mask, + output_attentions=output_attentions, + ) + attention_output = self_attention_outputs[0] + outputs = self_attention_outputs[1:] # add self attentions if we output attention weights + + # first residual connection + hidden_states = attention_output + residual + + # in Pix2StructVision, layernorm is also applied after self-attention + layer_output = self.pre_mlp_layer_norm(hidden_states) + layer_output = self.mlp(layer_output) + hidden_states # second residual connection + + outputs = (layer_output,) + outputs + + return outputs + + +class Pix2StructVisionEncoder(nn.Module): + def __init__(self, config: Pix2StructConfig) -> None: + super().__init__() + self.config = config + self.layer = nn.ModuleList([Pix2StructVisionLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ) -> Union[tuple, BaseModelOutput]: + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer_head_mask = head_mask[i] if head_mask is not None else None + + if self.gradient_checkpointing and self.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs, output_attentions) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(layer_module), + hidden_states, + attention_mask, + layer_head_mask, + ) + else: + layer_outputs = layer_module(hidden_states, attention_mask, layer_head_mask, output_attentions) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +class Pix2StructPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = Pix2StructConfig + + @property + def dummy_inputs(self): + input_ids = torch.tensor(DUMMY_INPUTS) + input_mask = torch.tensor(DUMMY_MASK) + dummy_inputs = { + "decoder_input_ids": input_ids, + "input_ids": input_ids, + "decoder_attention_mask": input_mask, + } + return dummy_inputs + + def _init_weights(self, module): + """Initialize the weights""" + factor = self.config.initializer_factor # Used for testing weights initialization + if isinstance(module, Pix2StructLayerNorm): + module.weight.data.fill_(factor * 1.0) + elif isinstance(module, Pix2StructTextDenseGatedActDense): + hidden_size = ( + self.config.text_config.hidden_size + if isinstance(self.config, Pix2StructConfig) + else self.config.hidden_size + ) + d_ff = self.config.text_config.d_ff if isinstance(self.config, Pix2StructConfig) else self.config.d_ff + + module.wi_0.weight.data.normal_(mean=0.0, std=factor * ((hidden_size) ** -0.5)) + if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None: + module.wi_0.bias.data.zero_() + module.wi_1.weight.data.normal_(mean=0.0, std=factor * ((hidden_size) ** -0.5)) + if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None: + module.wi_1.bias.data.zero_() + module.wo.weight.data.normal_(mean=0.0, std=factor * ((d_ff) ** -0.5)) + if hasattr(module.wo, "bias") and module.wo.bias is not None: + module.wo.bias.data.zero_() + elif isinstance(module, Pix2StructTextAttention): + # Mesh TensorFlow attention initialization to avoid scaling before softmax + # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/attention.py#L136 + hidden_size = ( + self.config.text_config.hidden_size + if isinstance(self.config, Pix2StructConfig) + else self.config.hidden_size + ) + key_value_proj_dim = ( + self.config.text_config.d_kv if isinstance(self.config, Pix2StructConfig) else self.config.hidden_size + ) + n_heads = ( + self.config.text_config.num_heads + if isinstance(self.config, Pix2StructConfig) + else self.config.num_heads + ) + + module.query.weight.data.normal_(mean=0.0, std=factor * ((hidden_size * key_value_proj_dim) ** -0.5)) + module.key.weight.data.normal_(mean=0.0, std=factor * (hidden_size**-0.5)) + module.value.weight.data.normal_(mean=0.0, std=factor * (hidden_size**-0.5)) + module.output.weight.data.normal_(mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5)) + if module.has_relative_attention_bias: + module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((hidden_size) ** -0.5)) + elif isinstance(module, nn.Embedding): + hidden_size = ( + self.config.text_config.hidden_size + if isinstance(self.config, Pix2StructConfig) + else self.config.hidden_size + ) + + module.weight.data.normal_(mean=0.0, std=factor * ((hidden_size) ** -0.5)) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, Pix2StructTextModel): + hidden_size = ( + self.config.text_config.hidden_size + if isinstance(self.config, Pix2StructConfig) + else self.config.hidden_size + ) + + module.lm_head.weight.data.normal_(mean=0.0, std=factor * ((hidden_size) ** -0.5)) + elif isinstance(module, (nn.Linear, nn.Conv2d)): + # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid + # `trunc_normal_cpu` not implemented in `half` issues + module.weight.data = nn.init.trunc_normal_( + module.weight.data.to(torch.float32), mean=0.0, std=self.config.initializer_range + ).to(module.weight.dtype) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, Pix2StructLayerNorm): + if module.weight is not None: + module.weight.data.fill_(1.0) + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + # Copied from transformers.models.t5.modeling_t5.T5PreTrainedModel._shift_right with T5->Pix2Struct + def _shift_right(self, input_ids): + decoder_start_token_id = self.config.decoder_start_token_id + pad_token_id = self.config.pad_token_id + + assert decoder_start_token_id is not None, ( + "self.model.config.decoder_start_token_id has to be defined. In Pix2Struct it is usually set to the pad_token_id." + " See Pix2Struct docs for more information" + ) + + # shift inputs to the right + if is_torch_fx_proxy(input_ids): + # Item assignment is not supported natively for proxies. + shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), decoder_start_token_id) + shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) + else: + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() + shifted_input_ids[..., 0] = decoder_start_token_id + + assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined." + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) + + return shifted_input_ids + + +PIX2STRUCT_VISION_START_DOCSTRING = r""" + This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it + as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and + behavior. + + Parameters: + config ([`Pix2StructConfig`]): Model configuration class with all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +PIX2STRUCT_VISION_INPUTS_DOCSTRING = r""" + Args: + flattened_patches (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels x patch_height x patch_width)`): + Flattened and padded pixel values. These values can be obtained using [`AutoImageProcessor`]. See + [`Pix2StructVisionImageProcessor.__call__`] for details. Check the [original + paper](https://arxiv.org/abs/2210.03347) (figure 5) for more details. + + attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding pixel values. Mask values selected in `[0, 1]`: + + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + 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. +""" + + +@add_start_docstrings( + "The bare Pix2StructVision Model transformer outputting raw hidden-states without any specific head on top.", + PIX2STRUCT_VISION_START_DOCSTRING, +) +class Pix2StructVisionModel(Pix2StructPreTrainedModel): + config_class = Pix2StructVisionConfig + main_input_name = "flattened_patches" + supports_gradient_checkpointing = True + _no_split_modules = ["Pix2StructVisionLayer"] + + def __init__(self, config: Pix2StructConfig): + super().__init__(config) + self.config = config + + self.embeddings = Pix2StructVisionEmbeddings(config) + self.encoder = Pix2StructVisionEncoder(config) + + self.layernorm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + # Initialize weights and apply final processing + self.post_init() + + def _set_gradient_checkpointing(self, module: Pix2StructVisionEncoder, value: bool = False) -> None: + if isinstance(module, Pix2StructVisionEncoder): + module.gradient_checkpointing = value + + def get_input_embeddings(self): + return self.embeddings.patch_projection + + def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None: + """ + Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base + class PreTrainedModel + """ + for layer, heads in heads_to_prune.items(): + self.encoder.layer[layer].attention.prune_heads(heads) + + @add_start_docstrings_to_model_forward(PIX2STRUCT_VISION_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=_CONFIG_FOR_DOC) + def forward( + self, + flattened_patches: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + Example: + + ```python + >>> import requests + >>> from PIL import Image + >>> from transformers import AutoProcessor, Pix2StructVisionModel + + >>> image_processor = AutoProcessor.from_pretrained("google/pix2struct-textcaps-base") + >>> model = Pix2StructVisionModel.from_pretrained("google/pix2struct-textcaps-base") + + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = image_processor(images=image, return_tensors="pt") + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + >>> list(last_hidden_states.shape) + [1, 2048, 768] + ``` + """ + 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 flattened_patches is None: + raise ValueError("You have to specify flattened_patches") + + if attention_mask is None: + # check where `flattened_patches` is not 0 + attention_mask = (flattened_patches.sum(dim=-1) != 0).float() + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] + # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embedding_output = self.embeddings(flattened_patches) + + encoder_outputs = self.encoder( + embedding_output, + attention_mask=attention_mask, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = encoder_outputs[0] + sequence_output = self.layernorm(sequence_output) + + if not return_dict: + head_outputs = (sequence_output,) + return head_outputs + encoder_outputs[1:] + + return BaseModelOutput( + last_hidden_state=sequence_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +# Copied from transformers.models.t5.modeling_t5.T5DenseGatedActDense with T5->Pix2StructText,d_model->hidden_size +class Pix2StructTextDenseGatedActDense(nn.Module): + def __init__(self, config: Pix2StructTextConfig): + super().__init__() + self.wi_0 = nn.Linear(config.hidden_size, config.d_ff, bias=False) + self.wi_1 = nn.Linear(config.hidden_size, config.d_ff, bias=False) + self.wo = nn.Linear(config.d_ff, config.hidden_size, bias=False) + self.dropout = nn.Dropout(config.dropout_rate) + self.act = ACT2FN[config.dense_act_fn] + + def forward(self, hidden_states): + hidden_gelu = self.act(self.wi_0(hidden_states)) + hidden_linear = self.wi_1(hidden_states) + hidden_states = hidden_gelu * hidden_linear + hidden_states = self.dropout(hidden_states) + + # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32. + # See https://github.com/huggingface/transformers/issues/20287 + # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None`` + if ( + isinstance(self.wo.weight, torch.Tensor) + and hidden_states.dtype != self.wo.weight.dtype + and self.wo.weight.dtype != torch.int8 + ): + hidden_states = hidden_states.to(self.wo.weight.dtype) + + hidden_states = self.wo(hidden_states) + return hidden_states + + +class Pix2StructTextLayerFF(nn.Module): + def __init__(self, config: Pix2StructTextConfig): + super().__init__() + self.DenseReluDense = Pix2StructTextDenseGatedActDense(config) + + self.layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + # Copied from transformers.models.t5.modeling_t5.T5LayerFF.forward + def forward(self, hidden_states): + forwarded_states = self.layer_norm(hidden_states) + forwarded_states = self.DenseReluDense(forwarded_states) + hidden_states = hidden_states + self.dropout(forwarded_states) + return hidden_states + + +class Pix2StructTextAttention(nn.Module): + def __init__(self, config: Pix2StructTextConfig, has_relative_attention_bias=False): + super().__init__() + self.has_relative_attention_bias = has_relative_attention_bias + self.relative_attention_num_buckets = config.relative_attention_num_buckets + self.relative_attention_max_distance = config.relative_attention_max_distance + self.hidden_size = config.hidden_size + self.key_value_proj_dim = config.d_kv + self.n_heads = config.num_heads + self.dropout = config.dropout_rate + self.inner_dim = self.n_heads * self.key_value_proj_dim + + # Mesh TensorFlow initialization to avoid scaling before softmax + self.query = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.key = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.value = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.output = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + + if self.has_relative_attention_bias: + self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads) + self.pruned_heads = set() + self.gradient_checkpointing = False + + @staticmethod + # Copied from transformers.models.t5.modeling_t5.T5Attention._relative_position_bucket + def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128): + """ + Adapted from Mesh Tensorflow: + https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 + + Translate relative position to a bucket number for relative attention. The relative position is defined as + memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to + position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for + small absolute relative_position and larger buckets for larger absolute relative_positions. All relative + positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. + This should allow for more graceful generalization to longer sequences than the model has been trained on + + Args: + relative_position: an int32 Tensor + bidirectional: a boolean - whether the attention is bidirectional + num_buckets: an integer + max_distance: an integer + + Returns: + a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) + """ + relative_buckets = 0 + if bidirectional: + num_buckets //= 2 + relative_buckets += (relative_position > 0).to(torch.long) * num_buckets + relative_position = torch.abs(relative_position) + else: + relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) + # now relative_position is in the range [0, inf) + + # half of the buckets are for exact increments in positions + max_exact = num_buckets // 2 + is_small = relative_position < max_exact + + # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance + relative_position_if_large = max_exact + ( + torch.log(relative_position.float() / max_exact) + / math.log(max_distance / max_exact) + * (num_buckets - max_exact) + ).to(torch.long) + relative_position_if_large = torch.min( + relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1) + ) + + relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) + return relative_buckets + + # Adapted from transformers.models.t5.modeling_t5.T5Attention.compute_bias + def compute_bias(self, query_length, key_length, device=None): + """Compute binned relative position bias""" + if device is None: + device = self.relative_attention_bias.weight.device + context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] + memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] + relative_position = memory_position - context_position # shape (query_length, key_length) + relative_position_bucket = self._relative_position_bucket( + relative_position, # shape (query_length, key_length) + bidirectional=False, + num_buckets=self.relative_attention_num_buckets, + max_distance=self.relative_attention_max_distance, + ) + values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads) + values = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length) + return values + + def forward( + self, + hidden_states, + mask=None, + key_value_states=None, + position_bias=None, + past_key_value=None, + layer_head_mask=None, + query_length=None, + use_cache=False, + output_attentions=False, + ): + """ + Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). + """ + # Input is (batch_size, seq_length, dim) + # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length) + # past_key_value[0] is (batch_size, n_heads, q_len - 1, dim_per_head) + batch_size, seq_length = hidden_states.shape[:2] + + real_seq_length = seq_length + + if past_key_value is not None: + if len(past_key_value) != 2: + raise ValueError( + f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states" + ) + real_seq_length += past_key_value[0].shape[2] if query_length is None else query_length + + key_length = real_seq_length if key_value_states is None else key_value_states.shape[1] + + def to_projection_shape(states): + """projection""" + return states.contiguous().view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) + + def project(hidden_states, proj_layer, key_value_states, past_key_value): + """projects hidden states correctly to key/query states""" + if key_value_states is None: + # self-attn + # (batch_size, n_heads, seq_length, dim_per_head) + hidden_states = to_projection_shape(proj_layer(hidden_states)) + elif past_key_value is None: + # cross-attn + # (batch_size, n_heads, seq_length, dim_per_head) + hidden_states = to_projection_shape(proj_layer(key_value_states)) + + if past_key_value is not None: + if key_value_states is None: + # self-attn + # (batch_size, n_heads, key_length, dim_per_head) + hidden_states = torch.cat([past_key_value, hidden_states], dim=2) + elif past_key_value.shape[2] != key_value_states.shape[1]: + # checking that the `sequence_length` of the `past_key_value` is the same as + # the provided `key_value_states` to support prefix tuning + # cross-attn + # (batch_size, n_heads, seq_length, dim_per_head) + hidden_states = to_projection_shape(proj_layer(key_value_states)) + else: + # cross-attn + hidden_states = past_key_value + return hidden_states + + # get query states + # (batch_size, n_heads, seq_length, dim_per_head) + query_states = to_projection_shape(self.query(hidden_states)) + + # get key/value states + key_states = project( + hidden_states, self.key, key_value_states, past_key_value[0] if past_key_value is not None else None + ) + value_states = project( + hidden_states, self.value, key_value_states, past_key_value[1] if past_key_value is not None else None + ) + + # compute scores + scores = torch.matmul( + query_states, key_states.transpose(3, 2) + ) # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 + + if position_bias is None: + if not self.has_relative_attention_bias: + position_bias = torch.zeros( + (1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype + ) + if self.gradient_checkpointing and self.training: + position_bias.requires_grad = True + else: + position_bias = self.compute_bias(real_seq_length, key_length, device=scores.device) + + # if key and values are already calculated + # we want only the last query position bias + if past_key_value is not None: + position_bias = position_bias[:, :, -hidden_states.size(1) :, :] + + if mask is not None: + position_bias = position_bias + mask # (batch_size, n_heads, seq_length, key_length) + + if self.pruned_heads: + mask = torch.ones(position_bias.shape[1]) + mask[list(self.pruned_heads)] = 0 + position_bias_masked = position_bias[:, mask.bool()] + else: + position_bias_masked = position_bias + + scores += position_bias_masked + # (batch_size, n_heads, seq_length, key_length) + attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores) + + # (batch_size, n_heads, seq_length, key_length) + attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + + # Mask heads if we want to + if layer_head_mask is not None: + attn_weights = attn_weights * layer_head_mask + + attn_output = torch.matmul(attn_weights, value_states) + # (batch_size, seq_length, dim) + attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim) + + attn_output = self.output(attn_output) + + present_key_value_state = (key_states, value_states) if use_cache else None + outputs = (attn_output,) + (present_key_value_state,) + (position_bias,) + + if output_attentions: + outputs = outputs + (attn_weights,) + return outputs + + +# Copied from transformers.models.t5.modeling_t5.T5LayerSelfAttention with T5LayerNorm->Pix2StructLayerNorm,T5Attention->Pix2StructTextAttention,self.SelfAttention->self.attention,config.d_model->config.hidden_size +class Pix2StructTextLayerSelfAttention(nn.Module): + def __init__(self, config, has_relative_attention_bias=False): + super().__init__() + self.attention = Pix2StructTextAttention(config, has_relative_attention_bias=has_relative_attention_bias) + self.layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward( + self, + hidden_states, + attention_mask=None, + position_bias=None, + layer_head_mask=None, + past_key_value=None, + use_cache=False, + output_attentions=False, + ): + normed_hidden_states = self.layer_norm(hidden_states) + attention_output = self.attention( + normed_hidden_states, + mask=attention_mask, + position_bias=position_bias, + layer_head_mask=layer_head_mask, + past_key_value=past_key_value, + use_cache=use_cache, + output_attentions=output_attentions, + ) + hidden_states = hidden_states + self.dropout(attention_output[0]) + outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them + return outputs + + +# Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5LayerNorm->Pix2StructLayerNorm,T5Attention->Pix2StructTextAttention,self.EncDecAttention->self.attention,config.d_model->config.hidden_size +class Pix2StructTextLayerCrossAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.attention = Pix2StructTextAttention(config, has_relative_attention_bias=False) + self.layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward( + self, + hidden_states, + key_value_states, + attention_mask=None, + position_bias=None, + layer_head_mask=None, + past_key_value=None, + use_cache=False, + query_length=None, + output_attentions=False, + ): + normed_hidden_states = self.layer_norm(hidden_states) + attention_output = self.attention( + normed_hidden_states, + mask=attention_mask, + key_value_states=key_value_states, + position_bias=position_bias, + layer_head_mask=layer_head_mask, + past_key_value=past_key_value, + use_cache=use_cache, + query_length=query_length, + output_attentions=output_attentions, + ) + layer_output = hidden_states + self.dropout(attention_output[0]) + outputs = (layer_output,) + attention_output[1:] # add attentions if we output them + return outputs + + +class Pix2StructTextBlock(nn.Module): + def __init__(self, config, has_relative_attention_bias=False): + super().__init__() + + self.self_attention = Pix2StructTextLayerSelfAttention( + config, has_relative_attention_bias=has_relative_attention_bias + ) + + self.encoder_decoder_attention = Pix2StructTextLayerCrossAttention(config) + + self.mlp = Pix2StructTextLayerFF(config) + + def forward( + self, + hidden_states, + attention_mask=None, + position_bias=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + encoder_decoder_position_bias=None, + layer_head_mask=None, + cross_attn_layer_head_mask=None, + past_key_value=None, + use_cache=False, + output_attentions=False, + return_dict=True, + ): + if past_key_value is not None: + expected_num_past_key_values = 2 if encoder_hidden_states is None else 4 + + if len(past_key_value) != expected_num_past_key_values: + raise ValueError( + f"There should be {expected_num_past_key_values} past states. " + f"{'2 (past / key) for cross attention. ' if expected_num_past_key_values == 4 else ''}" + f"Got {len(past_key_value)} past key / value states" + ) + + self_attn_past_key_value = past_key_value[:2] + cross_attn_past_key_value = past_key_value[2:] + else: + self_attn_past_key_value, cross_attn_past_key_value = None, None + + self_attention_outputs = self.self_attention( + hidden_states, + attention_mask=attention_mask, + position_bias=position_bias, + layer_head_mask=layer_head_mask, + past_key_value=self_attn_past_key_value, + use_cache=use_cache, + output_attentions=output_attentions, + ) + hidden_states, present_key_value_state = self_attention_outputs[:2] + attention_outputs = self_attention_outputs[2:] # Keep self-attention outputs and relative position weights + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any(): + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + do_cross_attention = encoder_hidden_states is not None + if do_cross_attention: + # the actual query length is unknown for cross attention + # if using past key value states. Need to inject it here + if present_key_value_state is not None: + query_length = present_key_value_state[0].shape[2] + else: + query_length = None + + cross_attention_outputs = self.encoder_decoder_attention( + hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + position_bias=encoder_decoder_position_bias, + layer_head_mask=cross_attn_layer_head_mask, + past_key_value=cross_attn_past_key_value, + query_length=query_length, + use_cache=use_cache, + output_attentions=output_attentions, + ) + hidden_states = cross_attention_outputs[0] + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any(): + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + # Combine self attn and cross attn key value states + if present_key_value_state is not None: + present_key_value_state = present_key_value_state + cross_attention_outputs[1] + + # Keep cross-attention outputs and relative position weights + attention_outputs = attention_outputs + cross_attention_outputs[2:] + + # Apply Feed Forward layer + hidden_states = self.mlp(hidden_states) + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any(): + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + outputs = (hidden_states,) + + if use_cache: + outputs = outputs + (present_key_value_state,) + attention_outputs + else: + outputs = outputs + attention_outputs + + return outputs + + +PIX2STRUCT_START_DOCSTRING = r""" + + The Pix2Struct model was proposed in [Pix2Struct: Screenshot Parsing as Pretraining for Visual Language + Understanding](https://arxiv.org/abs/2210.03347) by Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, + Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, Kristina Toutanova. It's an encoder decoder + transformer pre-trained in a image-to-text setting. + + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config (Union[`Pix2StructConfig`, `Pix2StructTextConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +PIX2STRUCT_TEXT_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Pix2StructText is a model with relative position + embeddings so you should be able to pad the inputs on both the right and the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for detail. + + [What are input IDs?](../glossary#input-ids) + + To know more on how to prepare `input_ids` for pretraining take a look a [Pix2StructText + Training](./t5#training). + attention_mask (`torch.FloatTensor` 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) + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + Pix2StructText uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If + `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + To know more on how to prepare `decoder_input_ids` for pretraining take a look at [Pix2StructText + Training](./t5#training). + decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules in the encoder. Mask values selected in `[0, + 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, + 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in + `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): + Tuple consists of (`last_hidden_state`, `optional`: *hidden_states*, `optional`: *attentions*) + `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` is a sequence of hidden states at + the output of the last layer of the encoder. Used in the cross-attention of the decoder. + past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention layers. Can be used to speed up 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. + decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded + representation. If `past_key_values` is used, optionally only the last `decoder_inputs_embeds` have to be + input (see `past_key_values`). This is useful if you want more control over how to convert + `decoder_input_ids` indices into associated vectors than the model's internal embedding lookup matrix. + + If `decoder_input_ids` and `decoder_inputs_embeds` are both unset, `decoder_inputs_embeds` takes the value + of `inputs_embeds`. + + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + +PIX2STRUCT_INPUTS_DOCSTRING = r""" + Args: + flattened_patches (`torch.FloatTensor` of shape `(batch_size, seq_length, hidden_size)`): + Flattened pixel patches. the `hidden_size` is obtained by the following formula: `hidden_size` = + `num_channels` * `patch_size` * `patch_size` + + The process of flattening the pixel patches is done by `Pix2StructProcessor`. + + attention_mask (`torch.FloatTensor` 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) + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + Pix2StructText uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If + `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + To know more on how to prepare `decoder_input_ids` for pretraining take a look at [Pix2StructText + Training](./t5#training). + decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules in the encoder. Mask values selected in `[0, + 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, + 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in + `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): + Tuple consists of (`last_hidden_state`, `optional`: *hidden_states*, `optional`: *attentions*) + `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` is a sequence of hidden states at + the output of the last layer of the encoder. Used in the cross-attention of the decoder. + past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention layers. Can be used to speed up 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)`. + decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded + representation. If `past_key_values` is used, optionally only the last `decoder_inputs_embeds` have to be + input (see `past_key_values`). This is useful if you want more control over how to convert + `decoder_input_ids` indices into associated vectors than the model's internal embedding lookup matrix. + + If `decoder_input_ids` and `decoder_inputs_embeds` are both unset, `decoder_inputs_embeds` takes the value + of `inputs_embeds`. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss for the decoder. + + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The standalone text decoder of Pix2Struct", + PIX2STRUCT_START_DOCSTRING, +) +class Pix2StructTextModel(Pix2StructPreTrainedModel): + config_class = Pix2StructTextConfig + _no_split_modules = ["Pix2StructTextBlock"] + supports_gradient_checkpointing = True + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (Pix2StructTextAttention, Pix2StructTextModel)): + module.gradient_checkpointing = value + + def __init__(self, config): + super().__init__(config) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + + self.layer = nn.ModuleList( + [Pix2StructTextBlock(config, has_relative_attention_bias=bool(i == 0)) for i in range(config.num_layers)] + ) + self.final_layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + self.gradient_checkpointing = False + + # Copied from transformers.models.t5.modeling_t5.T5PreTrainedModel._reorder_cache + def _reorder_cache(self, past_key_values, beam_idx): + # if decoder past is not included in output + # speedy decoding is disabled and no need to reorder + if past_key_values is None: + logger.warning("You might want to consider setting `use_cache=True` to speed up decoding") + return past_key_values + + reordered_decoder_past = () + for layer_past_states in past_key_values: + # get the correct batch idx from layer past batch dim + # batch dim of `past` is at 2nd position + reordered_layer_past_states = () + for layer_past_state in layer_past_states: + # need to set correct `past` for each of the four key / value states + reordered_layer_past_states = reordered_layer_past_states + ( + layer_past_state.index_select(0, beam_idx.to(layer_past_state.device)), + ) + + assert reordered_layer_past_states[0].shape == layer_past_states[0].shape + assert len(reordered_layer_past_states) == len(layer_past_states) + + reordered_decoder_past = reordered_decoder_past + (reordered_layer_past_states,) + return reordered_decoder_past + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, new_embeddings): + self.embed_tokens = new_embeddings + + @add_start_docstrings_to_model_forward(PIX2STRUCT_TEXT_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=CausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + inputs_embeds=None, + head_mask=None, + cross_attn_head_mask=None, + past_key_values=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + labels=None, + return_dict=None, + **kwargs, + ): + r""" + Returns: + + Example: + + ```python + >>> from transformers import AutoProcessor, Pix2StructTextModel + + >>> processor = AutoProcessor.from_pretrained("google/pix2struct-textcaps-base") + >>> model = Pix2StructTextModel.from_pretrained("google/pix2struct-textcaps-base") + + >>> inputs = processor(text="Hello, my dog is cute", return_tensors="pt") + >>> outputs = model(**inputs) + >>> loss = outputs.loss + ``` + """ + use_cache = use_cache if use_cache is not None else self.config.use_cache + 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 not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + if inputs_embeds is None: + assert self.embed_tokens is not None, "You have to initialize the model with valid token embeddings" + inputs_embeds = self.embed_tokens(input_ids) + + batch_size, seq_length = input_shape + + # required mask seq length can be calculated via length of past + mask_seq_length = past_key_values[0][0].shape[2] + seq_length if past_key_values is not None else seq_length + + if attention_mask is None: + attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) + if encoder_attention_mask is None and encoder_hidden_states is not None: + encoder_seq_length = encoder_hidden_states.shape[1] + encoder_attention_mask = torch.ones( + batch_size, encoder_seq_length, device=inputs_embeds.device, dtype=torch.long + ) + + # initialize past_key_values with `None` if past does not exist + if past_key_values is None: + past_key_values = [None] * len(self.layer) + + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape) + + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if encoder_hidden_states is not None: + encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones(encoder_hidden_shape, device=inputs_embeds.device) + encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) + else: + encoder_extended_attention_mask = None + + # Prepare head mask if needed + head_mask = self.get_head_mask(head_mask, self.config.num_layers) + cross_attn_head_mask = self.get_head_mask(cross_attn_head_mask, self.config.num_layers) + present_key_value_states = () if use_cache else None + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + all_cross_attentions = () if (output_attentions) else None + position_bias = None + encoder_decoder_position_bias = None + + hidden_states = self.dropout(inputs_embeds) + + for i, (layer_module, past_key_value) in enumerate(zip(self.layer, past_key_values)): + layer_head_mask = head_mask[i] + cross_attn_layer_head_mask = cross_attn_head_mask[i] + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + def create_custom_forward(module): + def custom_forward(*inputs): + return tuple(module(*inputs, use_cache, output_attentions)) + + return custom_forward + + layer_outputs = checkpoint( + create_custom_forward(layer_module), + hidden_states, + extended_attention_mask, + position_bias, + encoder_hidden_states, + encoder_extended_attention_mask, + encoder_decoder_position_bias, + layer_head_mask, + cross_attn_layer_head_mask, + None, # past_key_value is always None with gradient checkpointing + ) + else: + layer_outputs = layer_module( + hidden_states, + attention_mask=extended_attention_mask, + position_bias=position_bias, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_extended_attention_mask, + encoder_decoder_position_bias=encoder_decoder_position_bias, + layer_head_mask=layer_head_mask, + cross_attn_layer_head_mask=cross_attn_layer_head_mask, + past_key_value=past_key_value, + use_cache=use_cache, + output_attentions=output_attentions, + ) + + # layer_outputs is a tuple with: + # hidden-states, key-value-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) + if use_cache is False: + layer_outputs = layer_outputs[:1] + (None,) + layer_outputs[1:] + + hidden_states, present_key_value_state = layer_outputs[:2] + + # We share the position biases between the layers - the first layer store them + # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights), + # (cross-attention position bias), (cross-attention weights) + position_bias = layer_outputs[2] + if encoder_hidden_states is not None: + encoder_decoder_position_bias = layer_outputs[4 if output_attentions else 3] + # append next layer key value states + if use_cache: + present_key_value_states = present_key_value_states + (present_key_value_state,) + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[2],) + all_cross_attentions = all_cross_attentions + (layer_outputs[3],) + + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.dropout(hidden_states) + + logits = self.lm_head(hidden_states) + + # Add last layer + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + loss = None + if labels is not None: + loss_fct = nn.CrossEntropyLoss(ignore_index=-100, reduction="mean", label_smoothing=0.1) + masked_labels = labels.masked_fill(labels == self.config.pad_token_id, -100) + + loss = loss_fct(logits.contiguous().view(-1, logits.size(-1)), masked_labels.contiguous().view(-1)) + + if not return_dict: + return tuple( + v + for v in [ + loss, + logits, + present_key_value_states, + all_hidden_states, + all_attentions, + all_cross_attentions, + ] + if v is not None + ) + return CausalLMOutputWithCrossAttentions( + loss=loss, + logits=logits, + past_key_values=present_key_value_states, + hidden_states=all_hidden_states, + attentions=all_attentions, + cross_attentions=all_cross_attentions, + ) + + def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **model_kwargs): + input_shape = input_ids.shape + # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly + if attention_mask is None: + attention_mask = input_ids.new_ones(input_shape) + + # cut decoder_input_ids if past_key_values is used + if past_key_values is not None: + input_ids = input_ids[:, -1:] + + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "encoder_hidden_states": model_kwargs.get("encoder_hidden_states", None), + "encoder_attention_mask": model_kwargs.get("encoder_attention_mask", None), + "is_decoder": True, + } + + +@add_start_docstrings( + "A conditional generation model with a language modeling head. Can be used for sequence generation tasks.", + PIX2STRUCT_START_DOCSTRING, +) +class Pix2StructForConditionalGeneration(Pix2StructPreTrainedModel): + config_class = Pix2StructConfig + main_input_name = "flattened_patches" + + _keys_to_ignore_on_load_missing = [ + r"encoder.embed_tokens.weight", + r"decoder.embed_tokens.weight", + ] + _keys_to_ignore_on_load_unexpected = [ + r"decoder.layer.0.layer.1.EncDecAttention.relative_attention_bias.weight", + ] + + def __init__(self, config: Pix2StructConfig): + super().__init__(config) + encoder_config = copy.deepcopy(config.vision_config) + self.encoder = Pix2StructVisionModel(encoder_config) + + decoder_config = copy.deepcopy(config.text_config) + self.decoder_start_token_id = decoder_config.pad_token_id + self.decoder_eos_token_ids = decoder_config.eos_token_id + self.decoder = Pix2StructTextModel(decoder_config) + + self.is_vqa = config.is_vqa + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.shared + + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.decoder.set_input_embeddings(new_embeddings) + + def get_decoder(self): + return self.decoder + + def get_encoder(self): + return self.encoder + + @add_start_docstrings_to_model_forward(PIX2STRUCT_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=Seq2SeqModelOutput, config_class=_CONFIG_FOR_DOC) + def forward( + self, + flattened_patches: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + decoder_attention_mask: Optional[torch.BoolTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + decoder_head_mask: Optional[torch.FloatTensor] = None, + cross_attn_head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + labels: Optional[torch.LongTensor] = None, + decoder_inputs_embeds: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.FloatTensor], Seq2SeqModelOutput]: + r""" + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, Pix2StructForConditionalGeneration + + >>> processor = AutoProcessor.from_pretrained("google/pix2struct-textcaps-base") + >>> model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-textcaps-base") + + >>> labels = "A stop sign is on the street corner." + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, text=labels, return_tensors="pt", add_special_tokens=True) + + >>> # forward pass + >>> outputs = model(**inputs) + >>> last_hidden_states = outputs.loss + ```""" + use_cache = use_cache if use_cache is not None else self.config.text_config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Encode if needed (training, first prediction pass) + if encoder_outputs is None: + encoder_outputs = self.encoder( + flattened_patches=flattened_patches, + attention_mask=attention_mask, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = encoder_outputs[0] + + if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None: + # get decoder inputs from shifting lm labels to the right + decoder_input_ids = self._shift_right(labels) + decoder_attention_mask = ( + decoder_attention_mask + if decoder_attention_mask is not None + else decoder_input_ids.ne(self.config.pad_token_id).float() + ) + # Always attend to the first token + decoder_attention_mask[:, 0] = 1 + + # Decode + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + inputs_embeds=decoder_inputs_embeds, + past_key_values=past_key_values, + encoder_hidden_states=hidden_states, + encoder_attention_mask=attention_mask, + head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + labels=labels, + return_dict=return_dict, + ) + + if not return_dict: + return decoder_outputs + encoder_outputs + + return Seq2SeqLMOutput( + loss=decoder_outputs.loss, + logits=decoder_outputs.logits, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) + + @torch.no_grad() + def generate( + self, + flattened_patches: torch.FloatTensor, + decoder_input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + decoder_attention_mask: Optional[torch.LongTensor] = None, + **generate_kwargs, + ): + r""" + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, Pix2StructForConditionalGeneration + + >>> processor = AutoProcessor.from_pretrained("google/pix2struct-textcaps-base") + >>> model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-textcaps-base") + + >>> conditional_text = "A stop sign" + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, text=conditional_text, return_tensors="pt", add_special_tokens=True) + + >>> # forward pass + >>> outputs = model.generate(**inputs) + >>> print(processor.batch_decode(outputs, skip_special_tokens=True)) + ['A stop sign the street with a sign that says yes'] + ```""" + batch_size, _, _ = flattened_patches.shape + + vision_outputs = self.encoder(flattened_patches=flattened_patches, attention_mask=attention_mask) + + image_embeds = vision_outputs[0] + + if isinstance(decoder_input_ids, torch.Tensor): + # check if the first element of `input_ids` is equal to `decoder_input_ids`: + if (decoder_input_ids[:, 0] != self.decoder_start_token_id).all().item(): + # add `decoder_input_ids` as first token to `input_ids` + decoder_input_ids = torch.cat( + [ + torch.ones((decoder_input_ids.shape[0], 1), dtype=torch.long, device=decoder_input_ids.device) + * self.decoder_start_token_id, + decoder_input_ids, + ], + dim=-1, + ) + + if decoder_attention_mask is not None: + decoder_attention_mask = torch.cat( + [ + torch.ones( + (decoder_attention_mask.shape[0], 1), + dtype=torch.long, + device=decoder_attention_mask.device, + ), + decoder_attention_mask, + ], + dim=-1, + ) + elif decoder_input_ids is None: + decoder_input_ids = ( + torch.LongTensor([[self.decoder_start_token_id]]).repeat(batch_size, 1).to(image_embeds.device) + ) + + if decoder_attention_mask is None: + decoder_attention_mask = torch.ones_like(decoder_input_ids).to(image_embeds.device) + + outputs = self.decoder.generate( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + encoder_hidden_states=image_embeds, + encoder_attention_mask=attention_mask, + **generate_kwargs, + ) + + return outputs diff --git a/src/transformers/models/pix2struct/processing_pix2struct.py b/src/transformers/models/pix2struct/processing_pix2struct.py new file mode 100644 index 0000000000..eaa9f0dc42 --- /dev/null +++ b/src/transformers/models/pix2struct/processing_pix2struct.py @@ -0,0 +1,162 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. +# +# 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. +""" +Processor class for Pix2Struct. +""" + +from typing import List, Optional, Union + +from ...processing_utils import ProcessorMixin +from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy +from ...utils import TensorType + + +class Pix2StructProcessor(ProcessorMixin): + r""" + Constructs a PIX2STRUCT processor which wraps a BERT tokenizer and PIX2STRUCT image processor into a single + processor. + + [`Pix2StructProcessor`] offers all the functionalities of [`Pix2StructImageProcessor`] and [`T5TokenizerFast`]. See + the docstring of [`~Pix2StructProcessor.__call__`] and [`~Pix2StructProcessor.decode`] for more information. + + Args: + image_processor (`Pix2StructImageProcessor`): + An instance of [`Pix2StructImageProcessor`]. The image processor is a required input. + tokenizer (Union[`T5TokenizerFast`, `T5Tokenizer`]): + An instance of ['T5TokenizerFast`] or ['T5Tokenizer`]. The tokenizer is a required input. + """ + attributes = ["image_processor", "tokenizer"] + image_processor_class = "Pix2StructImageProcessor" + tokenizer_class = ("T5Tokenizer", "T5TokenizerFast") + + def __init__(self, image_processor, tokenizer): + tokenizer.return_token_type_ids = False + super().__init__(image_processor, tokenizer) + + def __call__( + self, + images=None, + text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None, + add_special_tokens: bool = False, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy] = None, + max_length: Optional[int] = None, + max_patches: Optional[int] = 2048, + stride: int = 0, + pad_to_multiple_of: Optional[int] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_token_type_ids: bool = False, + return_length: bool = False, + verbose: bool = True, + return_tensors: Optional[Union[str, TensorType]] = None, + **kwargs, + ) -> BatchEncoding: + """ + This method uses [`Pix2StructImageProcessor.preprocess`] method to prepare image(s) for the model, and + [`T5TokenizerFast.__call__`] to prepare text for the model. + + Please refer to the docstring of the above two methods for more information. + """ + if images is None and text is None: + raise ValueError("You have to specify either images or text.") + + # Get only text + if images is None and not self.image_processor.is_vqa: + self.current_processor = self.tokenizer + text_encoding = self.tokenizer( + text=text, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_token_type_ids=return_token_type_ids, + return_length=return_length, + verbose=verbose, + return_tensors=return_tensors, + **kwargs, + ) + return text_encoding + + if not self.image_processor.is_vqa: + # add pixel_values + encoding_image_processor = self.image_processor( + images, return_tensors=return_tensors, max_patches=max_patches, **kwargs + ) + else: + # add pixel_values and bbox + encoding_image_processor = self.image_processor( + images, return_tensors=return_tensors, max_patches=max_patches, header_text=text, **kwargs + ) + + if text is not None and not self.image_processor.is_vqa: + text_encoding = self.tokenizer( + text=text, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_token_type_ids=return_token_type_ids, + return_length=return_length, + verbose=verbose, + return_tensors=return_tensors, + **kwargs, + ) + + if "attention_mask" in text_encoding: + text_encoding["decoder_attention_mask"] = text_encoding.pop("attention_mask") + if "input_ids" in text_encoding: + text_encoding["decoder_input_ids"] = text_encoding.pop("input_ids") + else: + text_encoding = None + + if text_encoding is not None: + encoding_image_processor.update(text_encoding) + + return encoding_image_processor + + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to Pix2StructTokenizerFast'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 Pix2StructTokenizerFast'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(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) diff --git a/src/transformers/utils/dummy_pt_objects.py b/src/transformers/utils/dummy_pt_objects.py index a80af49e27..b062a3bec4 100644 --- a/src/transformers/utils/dummy_pt_objects.py +++ b/src/transformers/utils/dummy_pt_objects.py @@ -5049,6 +5049,37 @@ class PerceiverPreTrainedModel(metaclass=DummyObject): requires_backends(self, ["torch"]) +PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST = None + + +class Pix2StructForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Pix2StructPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Pix2StructTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Pix2StructVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + PLBART_PRETRAINED_MODEL_ARCHIVE_LIST = None diff --git a/src/transformers/utils/dummy_vision_objects.py b/src/transformers/utils/dummy_vision_objects.py index 5960f1e4b4..ea83c6b920 100644 --- a/src/transformers/utils/dummy_vision_objects.py +++ b/src/transformers/utils/dummy_vision_objects.py @@ -387,6 +387,13 @@ class PerceiverImageProcessor(metaclass=DummyObject): requires_backends(self, ["vision"]) +class Pix2StructImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + class PoolFormerFeatureExtractor(metaclass=DummyObject): _backends = ["vision"] diff --git a/tests/models/pix2struct/__init__.py b/tests/models/pix2struct/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/models/pix2struct/test_image_processing_pix2struct.py b/tests/models/pix2struct/test_image_processing_pix2struct.py new file mode 100644 index 0000000000..e13e38b71c --- /dev/null +++ b/tests/models/pix2struct/test_image_processing_pix2struct.py @@ -0,0 +1,289 @@ +# coding=utf-8 +# Copyright 2023 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 + +import numpy as np +import requests + +from transformers.testing_utils import require_torch, require_vision +from transformers.utils import is_torch_available, is_vision_available + +from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs + + +if is_torch_available(): + import torch + +if is_vision_available(): + from PIL import Image + + from transformers import Pix2StructImageProcessor + + +class Pix2StructImageProcessingTester(unittest.TestCase): + def __init__( + self, + parent, + batch_size=7, + num_channels=3, + image_size=18, + min_resolution=30, + max_resolution=400, + size=None, + do_normalize=True, + do_convert_rgb=True, + patch_size=None, + ): + 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.size = size + self.do_normalize = do_normalize + self.do_convert_rgb = do_convert_rgb + self.max_patches = [512, 1024, 2048, 4096] + self.patch_size = patch_size if patch_size is not None else {"height": 16, "width": 16} + + def prepare_image_processor_dict(self): + return {"do_normalize": self.do_normalize, "do_convert_rgb": self.do_convert_rgb} + + def prepare_dummy_image(self): + img_url = "https://www.ilankelman.org/stopsigns/australia.jpg" + raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") + return raw_image + + +@require_torch +@require_vision +class Pix2StructImageProcessingTest(ImageProcessingSavingTestMixin, unittest.TestCase): + image_processing_class = Pix2StructImageProcessor if is_vision_available() else None + + def setUp(self): + self.image_processor_tester = Pix2StructImageProcessingTester(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_normalize")) + self.assertTrue(hasattr(image_processor, "do_convert_rgb")) + + def test_expected_patches(self): + dummy_image = self.image_processor_tester.prepare_dummy_image() + + image_processor = self.image_processing_class(**self.image_processor_dict) + max_patch = 2048 + + inputs = image_processor(dummy_image, return_tensors="pt", max_patches=max_patch) + self.assertTrue(torch.allclose(inputs.flattened_patches.mean(), torch.tensor(0.0606), atol=1e-3, rtol=1e-3)) + + def test_call_pil(self): + # Initialize image_processor + image_processor = self.image_processing_class(**self.image_processor_dict) + # create random PIL images + image_inputs = prepare_image_inputs(self.image_processor_tester, equal_resolution=False) + for image in image_inputs: + self.assertIsInstance(image, Image.Image) + + # Test not batched input + expected_hidden_dim = ( + (self.image_processor_tester.patch_size["height"] * self.image_processor_tester.patch_size["width"]) + * self.image_processor_tester.num_channels + ) + 2 + + for max_patch in self.image_processor_tester.max_patches: + # Test not batched input + encoded_images = image_processor( + image_inputs[0], return_tensors="pt", max_patches=max_patch + ).flattened_patches + self.assertEqual( + encoded_images.shape, + (1, max_patch, expected_hidden_dim), + ) + + # Test batched + encoded_images = image_processor( + image_inputs, return_tensors="pt", max_patches=max_patch + ).flattened_patches + self.assertEqual( + encoded_images.shape, + (self.image_processor_tester.batch_size, max_patch, expected_hidden_dim), + ) + + def test_call_vqa(self): + # Initialize image_processor + image_processor = self.image_processing_class(**self.image_processor_dict) + # create random PIL images + image_inputs = prepare_image_inputs(self.image_processor_tester, equal_resolution=False) + for image in image_inputs: + self.assertIsInstance(image, Image.Image) + + # Test not batched input + expected_hidden_dim = ( + (self.image_processor_tester.patch_size["height"] * self.image_processor_tester.patch_size["width"]) + * self.image_processor_tester.num_channels + ) + 2 + + image_processor.is_vqa = True + + for max_patch in self.image_processor_tester.max_patches: + # Test not batched input + with self.assertRaises(ValueError): + encoded_images = image_processor( + image_inputs[0], return_tensors="pt", max_patches=max_patch + ).flattened_patches + + dummy_text = "Hello" + + encoded_images = image_processor( + image_inputs[0], return_tensors="pt", max_patches=max_patch, header_text=dummy_text + ).flattened_patches + self.assertEqual( + encoded_images.shape, + (1, max_patch, expected_hidden_dim), + ) + + # Test batched + encoded_images = image_processor( + image_inputs, return_tensors="pt", max_patches=max_patch, header_text=dummy_text + ).flattened_patches + self.assertEqual( + encoded_images.shape, + (self.image_processor_tester.batch_size, max_patch, expected_hidden_dim), + ) + + def test_call_numpy(self): + # Initialize image_processor + image_processor = self.image_processing_class(**self.image_processor_dict) + # create random numpy tensors + image_inputs = prepare_image_inputs(self.image_processor_tester, equal_resolution=False, numpify=True) + for image in image_inputs: + self.assertIsInstance(image, np.ndarray) + + expected_hidden_dim = ( + (self.image_processor_tester.patch_size["height"] * self.image_processor_tester.patch_size["width"]) + * self.image_processor_tester.num_channels + ) + 2 + + for max_patch in self.image_processor_tester.max_patches: + # Test not batched input + encoded_images = image_processor( + image_inputs[0], return_tensors="pt", max_patches=max_patch + ).flattened_patches + self.assertEqual( + encoded_images.shape, + (1, max_patch, expected_hidden_dim), + ) + + # Test batched + encoded_images = image_processor( + image_inputs, return_tensors="pt", max_patches=max_patch + ).flattened_patches + self.assertEqual( + encoded_images.shape, + (self.image_processor_tester.batch_size, max_patch, expected_hidden_dim), + ) + + def test_call_pytorch(self): + # Initialize image_processor + image_processor = self.image_processing_class(**self.image_processor_dict) + # create random PyTorch tensors + image_inputs = prepare_image_inputs(self.image_processor_tester, equal_resolution=False, torchify=True) + for image in image_inputs: + self.assertIsInstance(image, torch.Tensor) + + # Test not batched input + expected_hidden_dim = ( + (self.image_processor_tester.patch_size["height"] * self.image_processor_tester.patch_size["width"]) + * self.image_processor_tester.num_channels + ) + 2 + + for max_patch in self.image_processor_tester.max_patches: + # Test not batched input + encoded_images = image_processor( + image_inputs[0], return_tensors="pt", max_patches=max_patch + ).flattened_patches + self.assertEqual( + encoded_images.shape, + (1, max_patch, expected_hidden_dim), + ) + + # Test batched + encoded_images = image_processor( + image_inputs, return_tensors="pt", max_patches=max_patch + ).flattened_patches + self.assertEqual( + encoded_images.shape, + (self.image_processor_tester.batch_size, max_patch, expected_hidden_dim), + ) + + +@require_torch +@require_vision +class Pix2StructImageProcessingTestFourChannels(ImageProcessingSavingTestMixin, unittest.TestCase): + image_processing_class = Pix2StructImageProcessor if is_vision_available() else None + + def setUp(self): + self.image_processor_tester = Pix2StructImageProcessingTester(self, num_channels=4) + self.expected_encoded_image_num_channels = 3 + + @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_normalize")) + self.assertTrue(hasattr(image_processor, "do_convert_rgb")) + + def test_call_pil_four_channels(self): + # Initialize image_processor + image_processor = self.image_processing_class(**self.image_processor_dict) + # create random PIL images + image_inputs = prepare_image_inputs(self.image_processor_tester, equal_resolution=False) + for image in image_inputs: + self.assertIsInstance(image, Image.Image) + + # Test not batched input + expected_hidden_dim = ( + (self.image_processor_tester.patch_size["height"] * self.image_processor_tester.patch_size["width"]) + * (self.image_processor_tester.num_channels - 1) + ) + 2 + + for max_patch in self.image_processor_tester.max_patches: + # Test not batched input + encoded_images = image_processor( + image_inputs[0], return_tensors="pt", max_patches=max_patch + ).flattened_patches + self.assertEqual( + encoded_images.shape, + (1, max_patch, expected_hidden_dim), + ) + + # Test batched + encoded_images = image_processor( + image_inputs, return_tensors="pt", max_patches=max_patch + ).flattened_patches + self.assertEqual( + encoded_images.shape, + (self.image_processor_tester.batch_size, max_patch, expected_hidden_dim), + ) diff --git a/tests/models/pix2struct/test_modeling_pix2struct.py b/tests/models/pix2struct/test_modeling_pix2struct.py new file mode 100644 index 0000000000..3896f8d845 --- /dev/null +++ b/tests/models/pix2struct/test_modeling_pix2struct.py @@ -0,0 +1,711 @@ +# coding=utf-8 +# Copyright 2023 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 Pix2Struct model. """ + + +import inspect +import os +import tempfile +import unittest + +import numpy as np +import requests + +from transformers import Pix2StructConfig, Pix2StructTextConfig, Pix2StructVisionConfig +from transformers.testing_utils import require_torch, require_vision, slow, torch_device +from transformers.utils import is_torch_available, is_vision_available + +from ...test_configuration_common import ConfigTester +from ...test_modeling_common import ( + ModelTesterMixin, + _config_zero_init, + floats_tensor, + ids_tensor, + random_attention_mask, +) + + +if is_torch_available(): + import torch + from torch import nn + + from transformers import ( + Pix2StructForConditionalGeneration, + Pix2StructProcessor, + Pix2StructTextModel, + Pix2StructVisionModel, + ) + from transformers.models.pix2struct.modeling_pix2struct import PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST + + +if is_vision_available(): + from PIL import Image + + +class Pix2StructVisionModelTester: + def __init__( + self, + parent, + batch_size=12, + image_size=30, + patch_size=2, + num_channels=3, + is_training=True, + hidden_size=12, + patch_embed_hidden_size=12, + projection_dim=32, + max_patches=64, + num_hidden_layers=5, + num_attention_heads=4, + intermediate_size=37, + dropout=0.1, + attention_dropout=0.1, + initializer_range=1e-10, + scope=None, + ): + self.parent = parent + self.batch_size = batch_size + self.image_size = image_size + self.patch_embed_hidden_size = patch_embed_hidden_size + self.patch_size = patch_size + self.num_channels = num_channels + self.is_training = is_training + self.hidden_size = hidden_size + self.max_patches = max_patches + self.seq_length = self.max_patches + self.patch_proj_dim = ((patch_size**2) * num_channels) + 2 + + self.projection_dim = projection_dim + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.intermediate_size = intermediate_size + self.dropout = dropout + self.attention_dropout = attention_dropout + self.initializer_range = initializer_range + self.scope = scope + + def prepare_config_and_inputs(self): + flattened_patches = floats_tensor([self.batch_size, self.max_patches, self.patch_proj_dim]) + config = self.get_config() + + return config, flattened_patches + + def get_config(self): + return Pix2StructVisionConfig( + image_size=self.image_size, + patch_size=self.patch_size, + num_channels=self.num_channels, + hidden_size=self.hidden_size, + projection_dim=self.projection_dim, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads, + intermediate_size=self.intermediate_size, + dropout=self.dropout, + attention_dropout=self.attention_dropout, + initializer_range=self.initializer_range, + patch_embed_hidden_size=self.patch_embed_hidden_size, + ) + + def create_and_check_model(self, config, flattened_patches): + model = Pix2StructVisionModel(config=config) + model.to(torch_device) + model.eval() + with torch.no_grad(): + result = model(flattened_patches) + self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size)) + + def prepare_config_and_inputs_for_common(self): + config_and_inputs = self.prepare_config_and_inputs() + config, flattened_patches = config_and_inputs + inputs_dict = { + "flattened_patches": flattened_patches, + "attention_mask": torch.randint(0, 2, (self.batch_size, self.max_patches)), + } + return config, inputs_dict + + +@require_torch +class Pix2StructVisionModelTest(ModelTesterMixin, unittest.TestCase): + """ + Here we also overwrite some of the tests of test_modeling_common.py, as Pix2Struct does not use input_ids, inputs_embeds, + attention_mask and seq_length. + """ + + all_model_classes = (Pix2StructVisionModel,) if is_torch_available() else () + fx_compatible = False + test_pruning = False + test_resize_embeddings = False + test_head_masking = False + + def setUp(self): + self.model_tester = Pix2StructVisionModelTester(self) + self.config_tester = ConfigTester( + self, config_class=Pix2StructVisionConfig, has_text_modality=False, hidden_size=37 + ) + + def test_config(self): + self.config_tester.run_common_tests() + + @unittest.skip(reason="Pix2StructVision does not use inputs_embeds") + def test_inputs_embeds(self): + pass + + def test_model_common_attributes(self): + config, _ = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config) + self.assertIsInstance(model.get_input_embeddings(), (nn.Module)) + x = model.get_output_embeddings() + self.assertTrue(x is None or isinstance(x, nn.Linear)) + + def test_forward_signature(self): + config, _ = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config) + signature = inspect.signature(model.forward) + # signature.parameters is an OrderedDict => so arg_names order is deterministic + arg_names = [*signature.parameters.keys()] + + expected_arg_names = ["flattened_patches"] + self.assertListEqual(arg_names[:1], expected_arg_names) + + def test_model(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_model(*config_and_inputs) + + @unittest.skip(reason="Training is tested directly on `Pix2StructTextImageModelTest`") + def test_training(self): + pass + + @unittest.skip(reason="Training is tested directly on `Pix2StructTextImageModelTest`") + def test_training_gradient_checkpointing(self): + pass + + @unittest.skip(reason="Training is tested directly on `Pix2StructTextImageModelTest`") + def test_retain_grad_hidden_states_attentions(self): + pass + + @unittest.skip(reason="Pix2StructVisionModel has no base class and is not available in MODEL_MAPPING") + def test_save_load_fast_init_from_base(self): + pass + + @unittest.skip(reason="Pix2StructVisionModel has no base class and is not available in MODEL_MAPPING") + def test_save_load_fast_init_to_base(self): + pass + + @slow + def test_model_from_pretrained(self): + for model_name in PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]: + model = Pix2StructVisionModel.from_pretrained(model_name) + self.assertIsNotNone(model) + + +class Pix2StructTextModelTester: + def __init__( + self, + parent, + batch_size=12, + seq_length=7, + is_training=True, + use_input_mask=True, + use_labels=True, + vocab_size=99, + hidden_size=12, + projection_dim=32, + num_hidden_layers=5, + num_attention_heads=4, + intermediate_size=37, + dropout=0.1, + attention_dropout=0.1, + max_position_embeddings=512, + initializer_range=0.02, + bos_token_id=0, + scope=None, + ): + self.parent = parent + self.batch_size = batch_size + self.seq_length = seq_length + self.is_training = is_training + self.use_input_mask = use_input_mask + self.use_labels = use_labels + self.d_kv = hidden_size // num_attention_heads + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.projection_dim = projection_dim + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.intermediate_size = intermediate_size + self.dropout = dropout + self.attention_dropout = attention_dropout + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.scope = scope + self.bos_token_id = bos_token_id + + def prepare_config_and_inputs(self): + input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) + + input_mask = None + if self.use_input_mask: + input_mask = random_attention_mask([self.batch_size, self.seq_length]) + + if input_mask is not None: + batch_size, seq_length = input_mask.shape + rnd_start_indices = np.random.randint(1, seq_length - 1, size=(batch_size,)) + for batch_idx, start_index in enumerate(rnd_start_indices): + input_mask[batch_idx, :start_index] = 1 + input_mask[batch_idx, start_index:] = 0 + + config = self.get_config() + + return config, input_ids, input_mask + + def get_config(self): + return Pix2StructTextConfig( + vocab_size=self.vocab_size, + hidden_size=self.hidden_size, + projection_dim=self.projection_dim, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads, + intermediate_size=self.intermediate_size, + dropout=self.dropout, + attention_dropout=self.attention_dropout, + max_position_embeddings=self.max_position_embeddings, + initializer_range=self.initializer_range, + bos_token_id=self.bos_token_id, + d_kv=self.d_kv, + ) + + def create_and_check_model(self, config, input_ids, input_mask): + model = Pix2StructTextModel(config=config) + model.to(torch_device) + model.eval() + with torch.no_grad(): + result = model(input_ids, attention_mask=input_mask) + result = model(input_ids) + self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size)) + + def prepare_config_and_inputs_for_common(self): + config_and_inputs = self.prepare_config_and_inputs() + config, input_ids, input_mask = config_and_inputs + inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask} + return config, inputs_dict + + +@require_torch +class Pix2StructTextModelTest(ModelTesterMixin, unittest.TestCase): + all_model_classes = (Pix2StructTextModel,) if is_torch_available() else () + fx_compatible = False + test_pruning = False + test_head_masking = False + + def setUp(self): + self.model_tester = Pix2StructTextModelTester(self) + self.config_tester = ConfigTester(self, config_class=Pix2StructTextConfig, hidden_size=37) + + def test_config(self): + self.config_tester.run_common_tests() + + def test_model(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_model(*config_and_inputs) + + @unittest.skip(reason="Training is tested directly on `Pix2StructTextImageModelTest`") + def test_training(self): + pass + + @unittest.skip(reason="Training is tested directly on `Pix2StructTextImageModelTest`") + def test_training_gradient_checkpointing(self): + pass + + @unittest.skip(reason="Pix2Struct does not use inputs_embeds") + def test_inputs_embeds(self): + pass + + @unittest.skip(reason="Pix2StructTextModel has no base class and is not available in MODEL_MAPPING") + def test_save_load_fast_init_from_base(self): + pass + + @unittest.skip(reason="Pix2StructTextModel has no base class and is not available in MODEL_MAPPING") + def test_save_load_fast_init_to_base(self): + pass + + @slow + def test_model_from_pretrained(self): + for model_name in PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]: + model = Pix2StructTextModel.from_pretrained(model_name) + self.assertIsNotNone(model) + + +class Pix2StructTextImageModelsModelTester: + def __init__(self, parent, text_kwargs=None, vision_kwargs=None, is_training=True): + if text_kwargs is None: + text_kwargs = {} + if vision_kwargs is None: + vision_kwargs = {} + + self.parent = parent + self.text_model_tester = Pix2StructTextModelTester(parent, **text_kwargs) + self.vision_model_tester = Pix2StructVisionModelTester(parent, **vision_kwargs) + self.is_training = is_training + + def prepare_config_and_inputs(self): + text_config, input_ids, attention_mask = self.text_model_tester.prepare_config_and_inputs() + vision_config, flattened_patches = self.vision_model_tester.prepare_config_and_inputs() + + config = self.get_config(text_config, vision_config) + + return config, input_ids, attention_mask, flattened_patches + + def get_config(self, text_config, vision_config): + return Pix2StructConfig.from_text_vision_configs(text_config, vision_config, projection_dim=64) + + def prepare_config_and_inputs_for_common(self): + config_and_inputs = self.prepare_config_and_inputs() + config, input_ids, decoder_attention_mask, flattened_patches = config_and_inputs + + attention_mask = (flattened_patches.sum(dim=-1) != 0).float() + + inputs_dict = { + "decoder_input_ids": input_ids, + "labels": input_ids, + "decoder_attention_mask": decoder_attention_mask, + "flattened_patches": flattened_patches, + "attention_mask": attention_mask, + } + return config, inputs_dict + + +@require_torch +class Pix2StructTextImageModelTest(ModelTesterMixin, unittest.TestCase): + all_model_classes = (Pix2StructForConditionalGeneration,) if is_torch_available() else () + fx_compatible = False + test_head_masking = False + test_pruning = False + test_resize_embeddings = False + test_attention_outputs = False + test_torchscript = False + + def setUp(self): + self.model_tester = Pix2StructTextImageModelsModelTester(self) + + def test_model(self): + config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() + for model_class in self.all_model_classes: + model = model_class(config) + + output = model(**input_dict) + self.assertEqual( + output[1].shape, + ( + self.model_tester.vision_model_tester.batch_size, + self.model_tester.text_model_tester.seq_length, + self.model_tester.text_model_tester.vocab_size, + ), + ) + + @unittest.skip(reason="Hidden_states is tested in individual model tests") + def test_hidden_states_output(self): + pass + + @unittest.skip(reason="Inputs_embeds is tested in individual model tests") + def test_inputs_embeds(self): + pass + + @unittest.skip(reason="Retain_grad is tested in individual model tests") + def test_retain_grad_hidden_states_attentions(self): + pass + + @unittest.skip(reason="Pix2StructModel does not have input/output embeddings") + def test_model_common_attributes(self): + pass + + def test_forward_signature(self): + config, _ = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config) + signature = inspect.signature(model.forward) + # signature.parameters is an OrderedDict => so arg_names order is deterministic + arg_names = [*signature.parameters.keys()] + + if model.config.is_encoder_decoder: + expected_arg_names = [ + "input_ids", + "attention_mask", + "decoder_input_ids", + "decoder_attention_mask", + ] + expected_arg_names.extend( + ["head_mask", "decoder_head_mask", "cross_attn_head_mask", "encoder_outputs"] + if "head_mask" and "decoder_head_mask" and "cross_attn_head_mask" in arg_names + else ["encoder_outputs"] + ) + self.assertListEqual(arg_names[: len(expected_arg_names)], expected_arg_names) + else: + expected_arg_names = ( + ["input_ids"] if model_class != Pix2StructForConditionalGeneration else ["flattened_patches"] + ) + self.assertListEqual(arg_names[:1], expected_arg_names) + + def test_training(self): + if not self.model_tester.is_training: + return + + for model_class in self.all_model_classes[:-1]: + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + config.return_dict = True + + model = model_class(config) + model.to(torch_device) + model.train() + inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True) + + # hardcode labels to be the same as input_ids + inputs["labels"] = inputs["input_ids"] + + loss = model(**inputs).loss + loss.backward() + + def test_training_gradient_checkpointing(self): + if not self.model_tester.is_training: + return + + for model_class in self.all_model_classes[:-1]: + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + config.use_cache = False + config.return_dict = True + + model = model_class(config) + model.to(torch_device) + model.gradient_checkpointing_enable() + model.train() + inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True) + + # hardcode labels to be the same as input_ids + inputs["labels"] = inputs["input_ids"] + + loss = model(**inputs).loss + loss.backward() + + # override as the `logit_scale` parameter initilization is different for Pix2Struct + 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: + # check if `logit_scale` is initilized as per the original implementation + if name == "logit_scale": + self.assertAlmostEqual( + param.data.item(), + np.log(1 / 0.07), + delta=1e-3, + msg=f"Parameter {name} of model {model_class} seems not properly initialized", + ) + else: + 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", + ) + + def _create_and_check_torchscript(self, config, inputs_dict): + if not self.test_torchscript: + return + + configs_no_init = _config_zero_init(config) # To be sure we have no Nan + configs_no_init.torchscript = True + configs_no_init.return_dict = False + for model_class in self.all_model_classes: + model = model_class(config=configs_no_init) + model.to(torch_device) + model.eval() + + try: + input_ids = inputs_dict["input_ids"] + flattened_patches = inputs_dict["flattened_patches"] # Pix2Struct needs flattened_patches + traced_model = torch.jit.trace(model, (input_ids, flattened_patches)) + except RuntimeError: + self.fail("Couldn't trace module.") + + with tempfile.TemporaryDirectory() as tmp_dir_name: + pt_file_name = os.path.join(tmp_dir_name, "traced_model.pt") + + try: + torch.jit.save(traced_model, pt_file_name) + except Exception: + self.fail("Couldn't save module.") + + try: + loaded_model = torch.jit.load(pt_file_name) + except Exception: + self.fail("Couldn't load module.") + + model.to(torch_device) + model.eval() + + loaded_model.to(torch_device) + loaded_model.eval() + + model_state_dict = model.state_dict() + loaded_model_state_dict = loaded_model.state_dict() + + self.assertEqual(set(model_state_dict.keys()), set(loaded_model_state_dict.keys())) + + models_equal = True + for layer_name, p1 in model_state_dict.items(): + p2 = loaded_model_state_dict[layer_name] + if p1.data.ne(p2.data).sum() > 0: + models_equal = False + + self.assertTrue(models_equal) + + def test_load_vision_text_config(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + # Save Pix2StructConfig and check if we can load Pix2StructVisionConfig from it + with tempfile.TemporaryDirectory() as tmp_dir_name: + config.save_pretrained(tmp_dir_name) + vision_config = Pix2StructVisionConfig.from_pretrained(tmp_dir_name) + self.assertDictEqual(config.vision_config.to_dict(), vision_config.to_dict()) + + # Save Pix2StructConfig and check if we can load Pix2StructTextConfig from it + with tempfile.TemporaryDirectory() as tmp_dir_name: + config.save_pretrained(tmp_dir_name) + text_config = Pix2StructTextConfig.from_pretrained(tmp_dir_name) + self.assertDictEqual(config.text_config.to_dict(), text_config.to_dict()) + + +# We will verify our results on an image of a stop sign +def prepare_img(): + url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/australia.jpg" + im = Image.open(requests.get(url, stream=True).raw) + return im + + +@require_vision +@require_torch +@slow +class Pix2StructIntegrationTest(unittest.TestCase): + def test_inference_image_captioning(self): + model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-textcaps-base").to(torch_device) + processor = Pix2StructProcessor.from_pretrained("google/pix2struct-textcaps-base") + image = prepare_img() + + # image only + inputs = processor(images=image, return_tensors="pt").to(torch_device) + + predictions = model.generate(**inputs) + + self.assertEqual( + processor.decode(predictions[0], skip_special_tokens=True), "A stop sign is on a street corner." + ) + + def test_batched_inference_image_captioning(self): + model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-textcaps-base").to(torch_device) + processor = Pix2StructProcessor.from_pretrained("google/pix2struct-textcaps-base") + image_1 = prepare_img() + + second_url = ( + "https://www.connollycove.com/wp-content/uploads/2019/06/temple-bar-dublin-world-famous-irish-pub.jpg" + ) + image_2 = Image.open(requests.get(second_url, stream=True).raw) + + # image only + inputs = processor(images=[image_1, image_2], return_tensors="pt").to(torch_device) + + predictions = model.generate(**inputs) + + self.assertEqual( + processor.decode(predictions[0], skip_special_tokens=True), "A stop sign is on a street corner." + ) + + self.assertEqual( + processor.decode(predictions[1], skip_special_tokens=True), + "A row of books including The Temple Bar and Guiness.", + ) + + def test_batched_inference_image_captioning_conditioned(self): + model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-textcaps-base").to(torch_device) + processor = Pix2StructProcessor.from_pretrained("google/pix2struct-textcaps-base") + image_1 = prepare_img() + + second_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/temple-bar-dublin-world-famous-irish-pub.jpg" + image_2 = Image.open(requests.get(second_url, stream=True).raw) + texts = ["A picture of", "An photography of"] + + # image only + inputs = processor(images=[image_1, image_2], text=texts, return_tensors="pt").to(torch_device) + + predictions = model.generate(**inputs) + + self.assertEqual( + processor.decode(predictions[0], skip_special_tokens=True), "A picture of a stop sign that says yes." + ) + + self.assertEqual( + processor.decode(predictions[1], skip_special_tokens=True), + "An photography of the Temple Bar and a few other places.", + ) + + def test_vqa_model(self): + model_id = "ybelkada/pix2struct-ai2d-base" + + image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/ai2d-demo.jpg" + image = Image.open(requests.get(image_url, stream=True).raw) + + model = Pix2StructForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16).to( + torch_device + ) + processor = Pix2StructProcessor.from_pretrained(model_id) + + # image only + text = "What does the label 15 represent? (1) lava (2) core (3) tunnel (4) ash cloud" + + inputs = processor(images=image, return_tensors="pt", text=text).to(torch_device, torch.bfloat16) + + predictions = model.generate(**inputs) + self.assertEqual(processor.decode(predictions[0], skip_special_tokens=True), "ash cloud") + + def test_vqa_model_batched(self): + model_id = "ybelkada/pix2struct-ai2d-base" + + image_urls = [ + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/ai2d-demo.jpg", + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/ai2d-demo-2.png", + ] + + images = [Image.open(requests.get(image_url, stream=True).raw) for image_url in image_urls] + + texts = [ + "What does the label 15 represent? (1) lava (2) core (3) tunnel (4) ash cloud", + "What is the producer in the diagram? (1) Phytoplankton (2) Zooplankton (3) Large fish (4) Small fish", + ] + + model = Pix2StructForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16).to( + torch_device + ) + processor = Pix2StructProcessor.from_pretrained(model_id) + + inputs = processor(images=images, return_tensors="pt", text=texts).to(torch_device, torch.bfloat16) + + predictions = model.generate(**inputs) + self.assertEqual(processor.decode(predictions[0], skip_special_tokens=True), "ash cloud") + self.assertEqual(processor.decode(predictions[1], skip_special_tokens=True), "Phytoplankton") diff --git a/tests/models/pix2struct/test_processor_pix2struct.py b/tests/models/pix2struct/test_processor_pix2struct.py new file mode 100644 index 0000000000..6a7387b6ff --- /dev/null +++ b/tests/models/pix2struct/test_processor_pix2struct.py @@ -0,0 +1,192 @@ +# Copyright 2023 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 + +import numpy as np +import pytest + +from transformers.testing_utils import require_torch, require_vision +from transformers.utils import is_vision_available + + +if is_vision_available(): + from PIL import Image + + from transformers import ( + AutoProcessor, + Pix2StructImageProcessor, + Pix2StructProcessor, + PreTrainedTokenizerFast, + T5Tokenizer, + ) + + +@require_vision +@require_torch +class Pix2StructProcessorTest(unittest.TestCase): + def setUp(self): + self.tmpdirname = tempfile.mkdtemp() + + image_processor = Pix2StructImageProcessor() + tokenizer = T5Tokenizer.from_pretrained("t5-small") + + processor = Pix2StructProcessor(image_processor, tokenizer) + + 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 prepare_image_inputs(self): + """ + This function prepares a list of random PIL images of the same fixed size. + """ + + image_inputs = [np.random.randint(255, size=(3, 30, 400), dtype=np.uint8)] + + image_inputs = [Image.fromarray(np.moveaxis(x, 0, -1)) for x in image_inputs] + + return image_inputs + + def test_save_load_pretrained_additional_features(self): + processor = Pix2StructProcessor(tokenizer=self.get_tokenizer(), image_processor=self.get_image_processor()) + processor.save_pretrained(self.tmpdirname) + + tokenizer_add_kwargs = self.get_tokenizer(bos_token="(BOS)", eos_token="(EOS)") + image_processor_add_kwargs = self.get_image_processor(do_normalize=False, padding_value=1.0) + + processor = Pix2StructProcessor.from_pretrained( + self.tmpdirname, bos_token="(BOS)", eos_token="(EOS)", do_normalize=False, padding_value=1.0 + ) + + self.assertEqual(processor.tokenizer.get_vocab(), tokenizer_add_kwargs.get_vocab()) + self.assertIsInstance(processor.tokenizer, PreTrainedTokenizerFast) + + self.assertEqual(processor.image_processor.to_json_string(), image_processor_add_kwargs.to_json_string()) + self.assertIsInstance(processor.image_processor, Pix2StructImageProcessor) + + def test_image_processor(self): + image_processor = self.get_image_processor() + tokenizer = self.get_tokenizer() + + processor = Pix2StructProcessor(tokenizer=tokenizer, image_processor=image_processor) + + image_input = self.prepare_image_inputs() + + input_feat_extract = image_processor(image_input, return_tensors="np") + input_processor = processor(images=image_input, return_tensors="np") + + for key in input_feat_extract.keys(): + self.assertAlmostEqual(input_feat_extract[key].sum(), input_processor[key].sum(), delta=1e-2) + + def test_tokenizer(self): + image_processor = self.get_image_processor() + tokenizer = self.get_tokenizer() + + processor = Pix2StructProcessor(tokenizer=tokenizer, image_processor=image_processor) + + input_str = "lower newer" + + encoded_processor = processor(text=input_str) + + encoded_tok = tokenizer(input_str, return_token_type_ids=False, add_special_tokens=False) + + for key in encoded_tok.keys(): + self.assertListEqual(encoded_tok[key], encoded_processor[key]) + + def test_processor(self): + image_processor = self.get_image_processor() + tokenizer = self.get_tokenizer() + + processor = Pix2StructProcessor(tokenizer=tokenizer, image_processor=image_processor) + + input_str = "lower newer" + image_input = self.prepare_image_inputs() + + inputs = processor(text=input_str, images=image_input) + + self.assertListEqual( + list(inputs.keys()), ["flattened_patches", "attention_mask", "decoder_attention_mask", "decoder_input_ids"] + ) + + # test if it raises when no input is passed + with pytest.raises(ValueError): + processor() + + def test_processor_max_patches(self): + image_processor = self.get_image_processor() + tokenizer = self.get_tokenizer() + + processor = Pix2StructProcessor(tokenizer=tokenizer, image_processor=image_processor) + + input_str = "lower newer" + image_input = self.prepare_image_inputs() + + inputs = processor(text=input_str, images=image_input) + + max_patches = [512, 1024, 2048, 4096] + expected_hidden_size = [770, 770, 770, 770] + # with text + for i, max_patch in enumerate(max_patches): + inputs = processor(text=input_str, images=image_input, max_patches=max_patch) + self.assertEqual(inputs["flattened_patches"][0].shape[0], max_patch) + self.assertEqual(inputs["flattened_patches"][0].shape[1], expected_hidden_size[i]) + + # without text input + for i, max_patch in enumerate(max_patches): + inputs = processor(images=image_input, max_patches=max_patch) + self.assertEqual(inputs["flattened_patches"][0].shape[0], max_patch) + self.assertEqual(inputs["flattened_patches"][0].shape[1], expected_hidden_size[i]) + + def test_tokenizer_decode(self): + image_processor = self.get_image_processor() + tokenizer = self.get_tokenizer() + + processor = Pix2StructProcessor(tokenizer=tokenizer, image_processor=image_processor) + + predicted_ids = [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]] + + decoded_processor = processor.batch_decode(predicted_ids) + decoded_tok = tokenizer.batch_decode(predicted_ids) + + self.assertListEqual(decoded_tok, decoded_processor) + + def test_model_input_names(self): + image_processor = self.get_image_processor() + tokenizer = self.get_tokenizer() + + processor = Pix2StructProcessor(tokenizer=tokenizer, image_processor=image_processor) + + input_str = "lower newer" + image_input = self.prepare_image_inputs() + + inputs = processor(text=input_str, images=image_input) + + # For now the processor supports only ["flattened_patches", "input_ids", "attention_mask", "decoder_attention_mask"] + self.assertListEqual( + list(inputs.keys()), ["flattened_patches", "attention_mask", "decoder_attention_mask", "decoder_input_ids"] + ) + + inputs = processor(text=input_str) + + # For now the processor supports only ["flattened_patches", "input_ids", "attention_mask", "decoder_attention_mask"] + self.assertListEqual(list(inputs.keys()), ["input_ids", "attention_mask"]) diff --git a/utils/check_repo.py b/utils/check_repo.py index 121993bc1e..bcdecf729e 100644 --- a/utils/check_repo.py +++ b/utils/check_repo.py @@ -257,6 +257,9 @@ IGNORE_NON_AUTO_CONFIGURED = PRIVATE_MODELS.copy() + [ "FlaxCLIPVisionModel", "FlaxWav2Vec2ForCTC", "DetrForSegmentation", + "Pix2StructVisionModel", + "Pix2StructTextModel", + "Pix2StructForConditionalGeneration", "ConditionalDetrForSegmentation", "DPRReader", "FlaubertForQuestionAnswering",