


"""simple docstring"""




import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase						=      logging.get_logger(__name__)

__UpperCamelCase						=      {
    '''google/pix2struct-textcaps-base''': (
        '''https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json'''
    ),
}
class 					lowerCAmelCase	(					lowerCamelCase_    ):




		'''simple docstring'''







		SCREAMING_SNAKE_CASE_						:     List[str]								=     """pix2struct_text_model"""
		SCREAMING_SNAKE_CASE_						:     Any								=     ["""past_key_values"""]
		SCREAMING_SNAKE_CASE_						:     List[str]								=     {
		    """hidden_size""": """hidden_size""",
		    """num_attention_heads""": """num_heads""",
		    """num_hidden_layers""": """num_layers""",
		}







		def __init__(	self						,       lowerCAmelCase__=50_244						,       lowerCAmelCase__=768						,       lowerCAmelCase__=64						,       lowerCAmelCase__=2_048						,       lowerCAmelCase__=12						,       lowerCAmelCase__=12						,       lowerCAmelCase__=32						,       lowerCAmelCase__=128						,       lowerCAmelCase__=0.1						,       lowerCAmelCase__=1e-6						,       lowerCAmelCase__=1.0						,       lowerCAmelCase__="gelu_new"						,       lowerCAmelCase__=0						,       lowerCAmelCase__=False						,       lowerCAmelCase__=0						,       lowerCAmelCase__=1						,       lowerCAmelCase__=False						,       lowerCAmelCase__=True						,       **lowerCAmelCase__						,       )     ->							List[Any]:
								SCREAMING_SNAKE_CASE													=					vocab_size
								SCREAMING_SNAKE_CASE													=					hidden_size
								SCREAMING_SNAKE_CASE													=					d_kv
								SCREAMING_SNAKE_CASE													=					d_ff
								SCREAMING_SNAKE_CASE													=					num_layers
								SCREAMING_SNAKE_CASE													=					num_heads
								SCREAMING_SNAKE_CASE													=					relative_attention_num_buckets
								SCREAMING_SNAKE_CASE													=					relative_attention_max_distance
								SCREAMING_SNAKE_CASE													=					dropout_rate
								SCREAMING_SNAKE_CASE													=					layer_norm_epsilon
								SCREAMING_SNAKE_CASE													=					initializer_factor
								SCREAMING_SNAKE_CASE													=					use_cache

								SCREAMING_SNAKE_CASE													=					eos_token_id
								SCREAMING_SNAKE_CASE													=					decoder_start_token_id

								# for backwards compatibility
								SCREAMING_SNAKE_CASE													=					dense_act_fn

								super().__init__(
								    pad_token_id=lowerCAmelCase__						,       eos_token_id=lowerCAmelCase__						,       decoder_start_token_id=lowerCAmelCase__						,       tie_word_embeddings=lowerCAmelCase__						,       is_decoder=lowerCAmelCase__						,       **lowerCAmelCase__						,       )







		@classmethod
		def 							__A						(	cls						,       lowerCAmelCase__						,       **lowerCAmelCase__				)     ->							"PretrainedConfig":
								cls._set_token_in_kwargs(lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE							,				SCREAMING_SNAKE_CASE													=					cls.get_config_dict(lowerCAmelCase__						,       **lowerCAmelCase__				)

								# get the text config dict if we are loading from Pix2StructConfig
								if config_dict.get('model_type'				) == "pix2struct":
														SCREAMING_SNAKE_CASE													=					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(lowerCAmelCase__						,       **lowerCAmelCase__				)





class 					lowerCAmelCase	(					lowerCamelCase_    ):




		'''simple docstring'''







		SCREAMING_SNAKE_CASE_						:     List[str]								=     """pix2struct_vision_model"""







		def __init__(	self						,       lowerCAmelCase__=768						,       lowerCAmelCase__=768						,       lowerCAmelCase__=2_048						,       lowerCAmelCase__=64						,       lowerCAmelCase__=12						,       lowerCAmelCase__=12						,       lowerCAmelCase__="gelu_new"						,       lowerCAmelCase__=1e-6						,       lowerCAmelCase__=0.0						,       lowerCAmelCase__=0.0						,       lowerCAmelCase__=1e-10						,       lowerCAmelCase__=1.0						,       lowerCAmelCase__=4_096						,       lowerCAmelCase__=32						,       lowerCAmelCase__=128						,       **lowerCAmelCase__						,       )     ->							Tuple:
								super().__init__(**lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					hidden_size
								SCREAMING_SNAKE_CASE													=					patch_embed_hidden_size
								SCREAMING_SNAKE_CASE													=					d_ff
								SCREAMING_SNAKE_CASE													=					dropout_rate
								SCREAMING_SNAKE_CASE													=					num_hidden_layers
								SCREAMING_SNAKE_CASE													=					num_attention_heads
								SCREAMING_SNAKE_CASE													=					initializer_range
								SCREAMING_SNAKE_CASE													=					initializer_factor
								SCREAMING_SNAKE_CASE													=					attention_dropout
								SCREAMING_SNAKE_CASE													=					layer_norm_eps
								SCREAMING_SNAKE_CASE													=					dense_act_fn
								SCREAMING_SNAKE_CASE													=					seq_len
								SCREAMING_SNAKE_CASE													=					relative_attention_num_buckets
								SCREAMING_SNAKE_CASE													=					relative_attention_max_distance
								SCREAMING_SNAKE_CASE													=					d_kv







		@classmethod
		def 							__A						(	cls						,       lowerCAmelCase__						,       **lowerCAmelCase__				)     ->							"PretrainedConfig":
								cls._set_token_in_kwargs(lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE							,				SCREAMING_SNAKE_CASE													=					cls.get_config_dict(lowerCAmelCase__						,       **lowerCAmelCase__				)

								# get the vision config dict if we are loading from Pix2StructConfig
								if config_dict.get('model_type'				) == "pix2struct":
														SCREAMING_SNAKE_CASE													=					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(lowerCAmelCase__						,       **lowerCAmelCase__				)





class 					lowerCAmelCase	(					lowerCamelCase_    ):




		'''simple docstring'''







		SCREAMING_SNAKE_CASE_						:     Optional[int]								=     """pix2struct"""
		SCREAMING_SNAKE_CASE_						:     str								=     True







		def __init__(	self						,       lowerCAmelCase__=None						,       lowerCAmelCase__=None						,       lowerCAmelCase__=1.0						,       lowerCAmelCase__=0.02						,       lowerCAmelCase__=False						,       lowerCAmelCase__=False						,       lowerCAmelCase__=True						,       **lowerCAmelCase__						,       )     ->							Dict:
								super().__init__(tie_word_embeddings=lowerCAmelCase__						,       is_encoder_decoder=lowerCAmelCase__						,       **lowerCAmelCase__				)

								if text_config is None:
														SCREAMING_SNAKE_CASE													=					{}
														logger.info('text_config is None. Initializing the Pix2StructTextConfig with default values.'				)

								if vision_config is None:
														SCREAMING_SNAKE_CASE													=					{}
														logger.info('vision_config is None. Initializing the Pix2StructVisionConfig with default values.'				)

								SCREAMING_SNAKE_CASE													=					PixaStructTextConfig(**lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					PixaStructVisionConfig(**lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					self.text_config.decoder_start_token_id
								SCREAMING_SNAKE_CASE													=					self.text_config.pad_token_id
								SCREAMING_SNAKE_CASE													=					self.text_config.eos_token_id

								SCREAMING_SNAKE_CASE													=					initializer_factor
								SCREAMING_SNAKE_CASE													=					initializer_range

								SCREAMING_SNAKE_CASE													=					self.initializer_range
								SCREAMING_SNAKE_CASE													=					self.initializer_range

								SCREAMING_SNAKE_CASE													=					is_vqa







		@classmethod
		def 							__A						(	cls						,       lowerCAmelCase__						,       lowerCAmelCase__						,       **lowerCAmelCase__				)     ->							Tuple:
								return cls(text_config=text_config.to_dict()						,       vision_config=vision_config.to_dict()						,       **lowerCAmelCase__				)







		def 							__A						(	self				)     ->							List[Any]:
								SCREAMING_SNAKE_CASE													=					copy.deepcopy(self.__dict__				)
								SCREAMING_SNAKE_CASE													=					self.text_config.to_dict()
								SCREAMING_SNAKE_CASE													=					self.vision_config.to_dict()
								SCREAMING_SNAKE_CASE													=					self.__class__.model_type
								return output






"""simple docstring"""




import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam
class 					lowerCAmelCase	(					datasets.BeamBasedBuilder    ):




		'''simple docstring'''







		def 							__A						(	self				)     ->							Dict:
								return datasets.DatasetInfo(
								    features=datasets.Features({'content': datasets.Value('string'				)}				)						,       supervised_keys=lowerCAmelCase__						,       )







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__				)     ->							str:
								return [datasets.SplitGenerator(name=datasets.Split.TRAIN						,       gen_kwargs={'examples': get_test_dummy_examples()}				)]







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__				)     ->							Optional[int]:
								import apache_beam as beam

								return pipeline | "Load Examples" >> beam.Create(lowerCAmelCase__				)





class 					lowerCAmelCase	(					datasets.BeamBasedBuilder    ):




		'''simple docstring'''







		def 							__A						(	self				)     ->							int:
								return datasets.DatasetInfo(
								    features=datasets.Features({'a': datasets.Sequence({'b': datasets.Value('string'				)}				)}				)						,       supervised_keys=lowerCAmelCase__						,       )







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__				)     ->							Dict:
								return [
								    datasets.SplitGenerator(name=datasets.Split.TRAIN						,       gen_kwargs={'examples': get_test_nested_examples()}				)
								]







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__				)     ->							Optional[Any]:
								import apache_beam as beam

								return pipeline | "Load Examples" >> beam.Create(lowerCAmelCase__				)





def     lowercase			()   ->						str:
						return [(i, {"content": content}) for i, content in enumerate(['foo', 'bar', 'foobar']  )]


def     lowercase			()   ->						Optional[Any]:
						return [(i, {"a": {"b": [content]}}) for i, content in enumerate(['foo', 'bar', 'foobar']  )]


class 					lowerCAmelCase	(					lowerCamelCase_    ):




		'''simple docstring'''







		@require_beam
		def 							__A						(	self				)     ->							Optional[int]:
								SCREAMING_SNAKE_CASE													=					len(get_test_dummy_examples()				)
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
														SCREAMING_SNAKE_CASE													=					DummyBeamDataset(cache_dir=lowerCAmelCase__						,       beam_runner='DirectRunner'				)
														builder.download_and_prepare()
														self.assertTrue(
														    os.path.exists(
														        os.path.join(lowerCAmelCase__						,       builder.name						,       'default'						,       '0.0.0'						,       F'{builder.name}-train.arrow'				)				)				)
														self.assertDictEqual(builder.info.features						,       datasets.Features({'content': datasets.Value('string'				)}				)				)
														SCREAMING_SNAKE_CASE													=					builder.as_dataset()
														self.assertEqual(dset['train'].num_rows						,       lowerCAmelCase__				)
														self.assertEqual(dset['train'].info.splits['train'].num_examples						,       lowerCAmelCase__				)
														self.assertDictEqual(dset['train'][0]						,       get_test_dummy_examples()[0][1]				)
														self.assertDictEqual(
														    dset['train'][expected_num_examples - 1]						,       get_test_dummy_examples()[expected_num_examples - 1][1]				)
														self.assertTrue(
														    os.path.exists(os.path.join(lowerCAmelCase__						,       builder.name						,       'default'						,       '0.0.0'						,       'dataset_info.json'				)				)				)
														del dset







		@require_beam
		def 							__A						(	self				)     ->							int:
								import apache_beam as beam

								SCREAMING_SNAKE_CASE													=					beam.io.parquetio.WriteToParquet

								SCREAMING_SNAKE_CASE													=					len(get_test_dummy_examples()				)
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
														SCREAMING_SNAKE_CASE													=					DummyBeamDataset(cache_dir=lowerCAmelCase__						,       beam_runner='DirectRunner'				)
														with patch('apache_beam.io.parquetio.WriteToParquet'				) as write_parquet_mock:
																				SCREAMING_SNAKE_CASE													=					partial(lowerCAmelCase__						,       num_shards=2				)
																				builder.download_and_prepare()
														self.assertTrue(
														    os.path.exists(
														        os.path.join(
														            lowerCAmelCase__						,       builder.name						,       'default'						,       '0.0.0'						,       F'{builder.name}-train-00000-of-00002.arrow'				)				)				)
														self.assertTrue(
														    os.path.exists(
														        os.path.join(
														            lowerCAmelCase__						,       builder.name						,       'default'						,       '0.0.0'						,       F'{builder.name}-train-00000-of-00002.arrow'				)				)				)
														self.assertDictEqual(builder.info.features						,       datasets.Features({'content': datasets.Value('string'				)}				)				)
														SCREAMING_SNAKE_CASE													=					builder.as_dataset()
														self.assertEqual(dset['train'].num_rows						,       lowerCAmelCase__				)
														self.assertEqual(dset['train'].info.splits['train'].num_examples						,       lowerCAmelCase__				)
														# Order is not preserved when sharding, so we just check that all the elements are there
														self.assertListEqual(sorted(dset['train']['content']				)						,       sorted(['foo', 'bar', 'foobar']				)				)
														self.assertTrue(
														    os.path.exists(os.path.join(lowerCAmelCase__						,       builder.name						,       'default'						,       '0.0.0'						,       'dataset_info.json'				)				)				)
														del dset







		@require_beam
		def 							__A						(	self				)     ->							Optional[Any]:
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
														SCREAMING_SNAKE_CASE													=					DummyBeamDataset(cache_dir=lowerCAmelCase__				)
														self.assertRaises(datasets.builder.MissingBeamOptions						,       builder.download_and_prepare				)







		@require_beam
		def 							__A						(	self				)     ->							int:
								SCREAMING_SNAKE_CASE													=					len(get_test_nested_examples()				)
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
														SCREAMING_SNAKE_CASE													=					NestedBeamDataset(cache_dir=lowerCAmelCase__						,       beam_runner='DirectRunner'				)
														builder.download_and_prepare()
														self.assertTrue(
														    os.path.exists(
														        os.path.join(lowerCAmelCase__						,       builder.name						,       'default'						,       '0.0.0'						,       F'{builder.name}-train.arrow'				)				)				)
														self.assertDictEqual(
														    builder.info.features						,       datasets.Features({'a': datasets.Sequence({'b': datasets.Value('string'				)}				)}				)				)
														SCREAMING_SNAKE_CASE													=					builder.as_dataset()
														self.assertEqual(dset['train'].num_rows						,       lowerCAmelCase__				)
														self.assertEqual(dset['train'].info.splits['train'].num_examples						,       lowerCAmelCase__				)
														self.assertDictEqual(dset['train'][0]						,       get_test_nested_examples()[0][1]				)
														self.assertDictEqual(
														    dset['train'][expected_num_examples - 1]						,       get_test_nested_examples()[expected_num_examples - 1][1]				)
														self.assertTrue(
														    os.path.exists(os.path.join(lowerCAmelCase__						,       builder.name						,       'default'						,       '0.0.0'						,       'dataset_info.json'				)				)				)
														del dset


