





"""simple docstring"""





import os
import posixpath
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union

import numpy as np
import pyarrow as pa

import datasets
from datasets.arrow_writer import ArrowWriter, ParquetWriter
from datasets.config import MAX_SHARD_SIZE
from datasets.filesystems import (
    is_remote_filesystem,
    rename,
)
from datasets.iterable_dataset import _BaseExamplesIterable
from datasets.utils.py_utils import convert_file_size_to_int


_UpperCamelCase							=      datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
							import pyspark
@dataclass
class 		SCREAMING_SNAKE_CASE_							(    datasets.BuilderConfig   ):
				"""simple docstring"""
				__snake_case   :       Optional[datasets.Features]	 =		None



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							"pyspark.sql.DataFrame"    ,						SCREAMING_SNAKE_CASE       :							List[int]    ,						):


							'''simple docstring'''







							import pyspark

							def generate_fn():
														__lowerCamelCase      :      int         =df.select('''*'''    ,						pyspark.sql.functions.spark_partition_id().alias('''part_id'''   )   )
														for partition_id in partition_order:
																					__lowerCamelCase      :      List[str]         =df_with_partition_id.select('''*'''   ).where(F'part_id = {partition_id}'   ).drop('''part_id'''   )
																					__lowerCamelCase      :      Optional[Any]         =partition_df.collect()
																					__lowerCamelCase      :      Dict         =0
																					for row in rows:
																												yield F'{partition_id}_{row_id}', row.asDict()
																												row_id += 1

							return generate_fn
class 		SCREAMING_SNAKE_CASE_							(    _BaseExamplesIterable   ):
				"""simple docstring"""




				def __init__(     self  :List[Any]						,    __lowercase  :"pyspark.sql.DataFrame"						,    __lowercase  :List[str]=None						,    ):
											__lowerCamelCase      :      Union[str, Any]         =df
											__lowerCamelCase      :      str         =partition_order or range(self.df.rdd.getNumPartitions()			)
											__lowerCamelCase      :      Dict         =_generate_iterable_examples(self.df						,    self.partition_order			)




				def __iter__(     self  :Optional[int]			):
											yield from self.generate_examples_fn()




				def 					__lowercase  (     self  :str						,    __lowercase  :np.random.Generator			):
											__lowerCamelCase      :      Dict         =list(range(self.df.rdd.getNumPartitions()			)			)
											generator.shuffle(__lowercase			)
											return SparkExamplesIterable(self.df						,    partition_order=__lowercase			)




				def 					__lowercase  (     self  :Optional[int]						,    __lowercase  :int						,    __lowercase  :int			):
											__lowerCamelCase      :      Optional[int]         =self.split_shard_indices_by_worker(__lowercase						,    __lowercase			)
											return SparkExamplesIterable(self.df						,    partition_order=__lowercase			)




				@property
				def 					__lowercase  (     self  :List[str]			):
											return len(self.partition_order			)

class 		SCREAMING_SNAKE_CASE_							(    datasets.DatasetBuilder   ):
				"""simple docstring"""
				__snake_case   :       List[Any]	 =		SparkConfig




				def __init__(     self  :Union[str, Any]						,    __lowercase  :"pyspark.sql.DataFrame"						,    __lowercase  :str = None						,    __lowercase  :str = None						,    **__lowercase  :List[Any]						,    ):
											import pyspark

											__lowerCamelCase      :      str         =pyspark.sql.SparkSession.builder.getOrCreate()
											__lowerCamelCase      :      int         =df
											__lowerCamelCase      :      Any         =working_dir

											super().__init__(
											    cache_dir=__lowercase						,    config_name=str(self.df.semanticHash()			)						,    **__lowercase						,    )




				def 					__lowercase  (     self  :Any			):
											# Returns the path of the created file.
											def create_cache_and_write_probe(__lowercase  :int			):
																		# makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
																		# already exist.
																		os.makedirs(self._cache_dir						,    exist_ok=__lowercase			)
																		__lowerCamelCase      :      Dict         =os.path.join(self._cache_dir						,    '''fs_test''' + uuid.uuida().hex			)
																		# Opening the file in append mode will create a new file unless it already exists, in which case it will not
																		# change the file contents.
																		open(__lowercase						,    '''a'''			)
																		return [probe_file]

											if self._spark.conf.get('''spark.master'''						,    ''''''			).startswith('''local'''			):
																		return

											# If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
											# accessible to the driver.
											# TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
											if self._cache_dir:
																		__lowerCamelCase      :      List[str]         =(
																		    self._spark.sparkContext.parallelize(range(1			)						,    1			).mapPartitions(__lowercase			).collect()
																		)
																		if os.path.isfile(probe[0]			):
																									return

											raise ValueError(
											    '''When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir'''			)




				def 					__lowercase  (     self  :Any			):
											return datasets.DatasetInfo(features=self.config.features			)




				def 					__lowercase  (     self  :Tuple						,    __lowercase  :datasets.download.download_manager.DownloadManager			):
											return [datasets.SplitGenerator(name=datasets.Split.TRAIN			)]




				def 					__lowercase  (     self  :Union[str, Any]						,    __lowercase  :List[str]			):
											import pyspark

											def get_arrow_batch_size(__lowercase  :Dict			):
																		for batch in it:
																									yield pa.RecordBatch.from_pydict({'''batch_bytes''': [batch.nbytes]}			)

											__lowerCamelCase      :      List[str]         =self.df.count()
											__lowerCamelCase      :      str         =df_num_rows if df_num_rows <= 100 else 100
											# Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
											__lowerCamelCase      :      Tuple         =(
											    self.df.limit(__lowercase			)
											    .repartition(1			)
											    .mapInArrow(__lowercase						,    '''batch_bytes: long'''			)
											    .agg(pyspark.sql.functions.sum('''batch_bytes'''			).alias('''sample_bytes'''			)			)
											    .collect()[0]
											    .sample_bytes
											    / sample_num_rows
											)
											__lowerCamelCase      :      List[str]         =approx_bytes_per_row * df_num_rows
											if approx_total_size > max_shard_size:
																		# Make sure there is at least one row per partition.
																		__lowerCamelCase      :      str         =min(__lowercase						,    int(approx_total_size / max_shard_size			)			)
																		__lowerCamelCase      :      List[Any]         =self.df.repartition(__lowercase			)




				def 					__lowercase  (     self  :Dict						,    __lowercase  :str						,    __lowercase  :str						,    __lowercase  :int						,    ):
											import pyspark

											__lowerCamelCase      :      Optional[Any]         =ParquetWriter if file_format == '''parquet''' else ArrowWriter
											__lowerCamelCase      :      int         =os.path.join(self._working_dir						,    os.path.basename(__lowercase			)			) if self._working_dir else fpath
											__lowerCamelCase      :      Any         =file_format == '''parquet'''

											# Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
											# pickling the SparkContext.
											__lowerCamelCase      :      int         =self.config.features
											__lowerCamelCase      :      Any         =self._writer_batch_size
											__lowerCamelCase      :      Dict         =self._fs.storage_options

											def write_arrow(__lowercase  :Optional[int]			):
																		# Within the same SparkContext, no two task attempts will share the same attempt ID.
																		__lowerCamelCase      :      int         =pyspark.TaskContext().taskAttemptId()
																		__lowerCamelCase      :      Optional[Any]         =next(__lowercase						,    __lowercase			)
																		if first_batch is None:
																									# Some partitions might not receive any data.
																									return pa.RecordBatch.from_arrays(
																									    [[task_id], [0], [0]]						,    names=['''task_id''', '''num_examples''', '''num_bytes''']						,    )
																		__lowerCamelCase      :      Any         =0
																		__lowerCamelCase      :      int         =writer_class(
																		    features=__lowercase						,    path=working_fpath.replace('''SSSSS'''						,    f'{shard_id:05d}'			).replace('''TTTTT'''						,    f'{task_id:05d}'			)						,    writer_batch_size=__lowercase						,    storage_options=__lowercase						,    embed_local_files=__lowercase						,    )
																		__lowerCamelCase      :      Union[str, Any]         =pa.Table.from_batches([first_batch]			)
																		writer.write_table(__lowercase			)
																		for batch in it:
																									if max_shard_size is not None and writer._num_bytes >= max_shard_size:
																																__lowerCamelCase    ,						__lowerCamelCase      :      str         =writer.finalize()
																																writer.close()
																																yield pa.RecordBatch.from_arrays(
																																    [[task_id], [num_examples], [num_bytes]]						,    names=['''task_id''', '''num_examples''', '''num_bytes''']						,    )
																																shard_id += 1
																																__lowerCamelCase      :      Tuple         =writer_class(
																																    features=writer._features						,    path=working_fpath.replace('''SSSSS'''						,    f'{shard_id:05d}'			).replace('''TTTTT'''						,    f'{task_id:05d}'			)						,    writer_batch_size=__lowercase						,    storage_options=__lowercase						,    embed_local_files=__lowercase						,    )
																									__lowerCamelCase      :      List[str]         =pa.Table.from_batches([batch]			)
																									writer.write_table(__lowercase			)

																		if writer._num_bytes > 0:
																									__lowerCamelCase    ,						__lowerCamelCase      :      Any         =writer.finalize()
																									writer.close()
																									yield pa.RecordBatch.from_arrays(
																									    [[task_id], [num_examples], [num_bytes]]						,    names=['''task_id''', '''num_examples''', '''num_bytes''']						,    )

																		if working_fpath != fpath:
																									for file in os.listdir(os.path.dirname(__lowercase			)			):
																																__lowerCamelCase      :      List[str]         =os.path.join(os.path.dirname(__lowercase			)						,    os.path.basename(__lowercase			)			)
																																shutil.move(__lowercase						,    __lowercase			)

											__lowerCamelCase      :      Optional[Any]         =(
											    self.df.mapInArrow(__lowercase						,    '''task_id: long, num_examples: long, num_bytes: long'''			)
											    .groupBy('''task_id'''			)
											    .agg(
											        pyspark.sql.functions.sum('''num_examples'''			).alias('''total_num_examples'''			)						,    pyspark.sql.functions.sum('''num_bytes'''			).alias('''total_num_bytes'''			)						,    pyspark.sql.functions.count('''num_bytes'''			).alias('''num_shards'''			)						,    pyspark.sql.functions.collect_list('''num_examples'''			).alias('''shard_lengths'''			)						,    )
											    .collect()
											)
											for row in stats:
																		yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)




				def 					__lowercase  (     self  :List[str]						,    __lowercase  :"datasets.SplitGenerator"						,    __lowercase  :str = "arrow"						,    __lowercase  :Optional[Union[str, int]] = None						,    __lowercase  :Optional[int] = None						,    **__lowercase  :Tuple						,    ):
											self._validate_cache_dir()

											__lowerCamelCase      :      str         =convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE			)
											self._repartition_df_if_needed(__lowercase			)
											__lowerCamelCase      :      Any         =not is_remote_filesystem(self._fs			)
											__lowerCamelCase      :      Tuple         =os.path.join if is_local else posixpath.join

											__lowerCamelCase      :      Optional[Any]         ='''-TTTTT-SSSSS-of-NNNNN'''
											__lowerCamelCase      :      Any         =f'{self.name}-{split_generator.name}{SUFFIX}.{file_format}'
											__lowerCamelCase      :      Any         =path_join(self._output_dir						,    __lowercase			)

											__lowerCamelCase      :      List[str]         =0
											__lowerCamelCase      :      List[str]         =0
											__lowerCamelCase      :      Optional[Any]         =0
											__lowerCamelCase      :      Union[str, Any]         =[]
											__lowerCamelCase      :      Union[str, Any]         =[]

											for task_id, content in self._prepare_split_single(__lowercase						,    __lowercase						,    __lowercase			):
																		(
																		    (
																		    __lowerCamelCase
																		)    ,						(
																		    __lowerCamelCase
																		)    ,						(
																		    __lowerCamelCase
																		)    ,						(
																		    __lowerCamelCase
																		)    ,						
																		)      :      Union[str, Any]         =content
																		if num_bytes > 0:
																									total_num_examples += num_examples
																									total_num_bytes += num_bytes
																									total_shards += num_shards
																									task_id_and_num_shards.append((task_id, num_shards)			)
																									all_shard_lengths.extend(__lowercase			)

											__lowerCamelCase      :      Optional[int]         =total_num_examples
											__lowerCamelCase      :      str         =total_num_bytes

											# should rename everything at the end
											logger.debug(f'Renaming {total_shards} shards.'			)
											if total_shards > 1:
																		__lowerCamelCase      :      Any         =all_shard_lengths

																		# Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
																		# pickling error due to pickling the SparkContext.
																		__lowerCamelCase      :      Dict         =self._fs

																		# use the -SSSSS-of-NNNNN pattern
																		def _rename_shard(
																		    __lowercase  :int						,    __lowercase  :int						,    __lowercase  :int						,    ):
																									rename(
																									    __lowercase						,    fpath.replace('''SSSSS'''						,    f'{shard_id:05d}'			).replace('''TTTTT'''						,    f'{task_id:05d}'			)						,    fpath.replace('''TTTTT-SSSSS'''						,    f'{global_shard_id:05d}'			).replace('''NNNNN'''						,    f'{total_shards:05d}'			)						,    )

																		__lowerCamelCase      :      Tuple         =[]
																		__lowerCamelCase      :      Optional[int]         =0
																		for i in range(len(__lowercase			)			):
																									__lowerCamelCase    ,						__lowerCamelCase      :      Tuple         =task_id_and_num_shards[i]
																									for shard_id in range(__lowercase			):
																																args.append([task_id, shard_id, global_shard_id]			)
																																global_shard_id += 1
																		self._spark.sparkContext.parallelize(__lowercase						,    len(__lowercase			)			).map(lambda __lowercase			: _rename_shard(*__lowercase			)			).collect()
											else:
																		# don't use any pattern
																		__lowerCamelCase      :      Tuple         =0
																		__lowerCamelCase      :      Optional[Any]         =task_id_and_num_shards[0][0]
																		self._rename(
																		    fpath.replace('''SSSSS'''						,    f'{shard_id:05d}'			).replace('''TTTTT'''						,    f'{task_id:05d}'			)						,    fpath.replace(__lowercase						,    ''''''			)						,    )




				def 					__lowercase  (     self  :Optional[Any]						,    __lowercase  :"datasets.SplitGenerator"						,    ):
											return SparkExamplesIterable(self.df			)











"""simple docstring"""





import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import (
    RobertaTokenizer,
    TrOCRConfig,
    TrOCRForCausalLM,
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    ViTConfig,
    ViTImageProcessor,
    ViTModel,
)
from transformers.utils import logging


logging.set_verbosity_info()
_UpperCamelCase							=      logging.get_logger(__name__)



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Union[str, Any]    ,						SCREAMING_SNAKE_CASE       :							List[str]   ):


							'''simple docstring'''







							__lowerCamelCase      :      Any         =[]
							for i in range(encoder_config.num_hidden_layers   ):
														# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.norm1.weight', F'encoder.encoder.layer.{i}.layernorm_before.weight')   )
														rename_keys.append((F'encoder.deit.blocks.{i}.norm1.bias', F'encoder.encoder.layer.{i}.layernorm_before.bias')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.attn.proj.weight', F'encoder.encoder.layer.{i}.attention.output.dense.weight')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.attn.proj.bias', F'encoder.encoder.layer.{i}.attention.output.dense.bias')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.norm2.weight', F'encoder.encoder.layer.{i}.layernorm_after.weight')   )
														rename_keys.append((F'encoder.deit.blocks.{i}.norm2.bias', F'encoder.encoder.layer.{i}.layernorm_after.bias')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.mlp.fc1.weight', F'encoder.encoder.layer.{i}.intermediate.dense.weight')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.mlp.fc1.bias', F'encoder.encoder.layer.{i}.intermediate.dense.bias')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.mlp.fc2.weight', F'encoder.encoder.layer.{i}.output.dense.weight')   )
														rename_keys.append((F'encoder.deit.blocks.{i}.mlp.fc2.bias', F'encoder.encoder.layer.{i}.output.dense.bias')   )

							# cls token, position embeddings and patch embeddings of encoder
							rename_keys.extend(
							    [
							        ('''encoder.deit.cls_token''', '''encoder.embeddings.cls_token'''),
							        ('''encoder.deit.pos_embed''', '''encoder.embeddings.position_embeddings'''),
							        ('''encoder.deit.patch_embed.proj.weight''', '''encoder.embeddings.patch_embeddings.projection.weight'''),
							        ('''encoder.deit.patch_embed.proj.bias''', '''encoder.embeddings.patch_embeddings.projection.bias'''),
							        ('''encoder.deit.norm.weight''', '''encoder.layernorm.weight'''),
							        ('''encoder.deit.norm.bias''', '''encoder.layernorm.bias'''),
							    ]   )

							return rename_keys



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							int    ,						SCREAMING_SNAKE_CASE       :							Optional[int]   ):


							'''simple docstring'''







							for i in range(encoder_config.num_hidden_layers   ):
														# queries, keys and values (only weights, no biases)
														__lowerCamelCase      :      str         =state_dict.pop(F'encoder.deit.blocks.{i}.attn.qkv.weight'   )

														__lowerCamelCase      :      Dict         =in_proj_weight[
														    : encoder_config.hidden_size, :
														]
														__lowerCamelCase      :      Any         =in_proj_weight[
														    encoder_config.hidden_size : encoder_config.hidden_size * 2, :
														]
														__lowerCamelCase      :      Dict         =in_proj_weight[
														    -encoder_config.hidden_size :, :
														]



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[Any]    ,						SCREAMING_SNAKE_CASE       :							Optional[int]    ,						SCREAMING_SNAKE_CASE       :							Dict   ):


							'''simple docstring'''







							__lowerCamelCase      :      Optional[int]         =dct.pop(SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Optional[Any]         =val



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[Any]   ):


							'''simple docstring'''







							if "handwritten" in checkpoint_url:
														__lowerCamelCase      :      List[Any]         ='''https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg'''  # industry
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-12.jpg" # have
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-10.jpg" # let
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"  #
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122.jpg"
							elif "printed" in checkpoint_url or "stage1" in checkpoint_url:
														__lowerCamelCase      :      List[Any]         ='''https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg'''
							__lowerCamelCase      :      int         =Image.open(requests.get(SCREAMING_SNAKE_CASE    ,						stream=SCREAMING_SNAKE_CASE   ).raw   ).convert('''RGB'''   )
							return im



@torch.no_grad()
def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[str]    ,						SCREAMING_SNAKE_CASE       :							Any   ):


							'''simple docstring'''







							__lowerCamelCase      :      Union[str, Any]         =ViTConfig(image_size=384    ,						qkv_bias=SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Optional[int]         =TrOCRConfig()

							# size of the architecture
							if "base" in checkpoint_url:
														__lowerCamelCase      :      Any         =768
							elif "large" in checkpoint_url:
														# use ViT-large encoder
														__lowerCamelCase      :      Optional[int]         =1024
														__lowerCamelCase      :      List[Any]         =4096
														__lowerCamelCase      :      Dict         =24
														__lowerCamelCase      :      Optional[Any]         =16
														__lowerCamelCase      :      Union[str, Any]         =1024
							else:
														raise ValueError('''Should either find \'base\' or \'large\' in checkpoint URL'''   )

							# the large-printed + stage1 checkpoints uses sinusoidal position embeddings, no layernorm afterwards
							if "large-printed" in checkpoint_url or "stage1" in checkpoint_url:
														__lowerCamelCase      :      Optional[int]         =False
														__lowerCamelCase      :      int         ='''relu'''
														__lowerCamelCase      :      Any         =1024
														__lowerCamelCase      :      str         =True
														__lowerCamelCase      :      List[Any]         =False
														__lowerCamelCase      :      Optional[int]         =False

							# load HuggingFace model
							__lowerCamelCase      :      Dict         =ViTModel(SCREAMING_SNAKE_CASE    ,						add_pooling_layer=SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Any         =TrOCRForCausalLM(SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      List[Any]         =VisionEncoderDecoderModel(encoder=SCREAMING_SNAKE_CASE    ,						decoder=SCREAMING_SNAKE_CASE   )
							model.eval()

							# load state_dict of original model, rename some keys
							__lowerCamelCase      :      Optional[Any]         =torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE    ,						map_location='''cpu'''    ,						check_hash=SCREAMING_SNAKE_CASE   )['''model''']

							__lowerCamelCase      :      Optional[int]         =create_rename_keys(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )
							for src, dest in rename_keys:
														rename_key(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )
							read_in_q_k_v(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )

							# remove parameters we don't need
							del state_dict["encoder.deit.head.weight"]
							del state_dict["encoder.deit.head.bias"]
							del state_dict["decoder.version"]

							# add prefix to decoder keys
							for key, val in state_dict.copy().items():
														__lowerCamelCase      :      List[Any]         =state_dict.pop(SCREAMING_SNAKE_CASE   )
														if key.startswith('''decoder'''   ) and "output_projection" not in key:
																					__lowerCamelCase      :      List[Any]         =val
														else:
																					__lowerCamelCase      :      Optional[Any]         =val

    # load state dict
							model.load_state_dict(SCREAMING_SNAKE_CASE   )

							# Check outputs on an image
							__lowerCamelCase      :      int         =ViTImageProcessor(size=encoder_config.image_size   )
							__lowerCamelCase      :      Dict         =RobertaTokenizer.from_pretrained('''roberta-large'''   )
							__lowerCamelCase      :      str         =TrOCRProcessor(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )

							__lowerCamelCase      :      str         =processor(images=prepare_img(SCREAMING_SNAKE_CASE   )    ,						return_tensors='''pt'''   ).pixel_values

							# verify logits
							__lowerCamelCase      :      List[Any]         =torch.tensor([[model.config.decoder.decoder_start_token_id]]   )
							__lowerCamelCase      :      Optional[int]         =model(pixel_values=SCREAMING_SNAKE_CASE    ,						decoder_input_ids=SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Optional[Any]         =outputs.logits

							__lowerCamelCase      :      Optional[Any]         =torch.Size([1, 1, 50265]   )
							if "trocr-base-handwritten" in checkpoint_url:
														__lowerCamelCase      :      List[str]         =torch.tensor(
														    [-1.4_502, -4.6_683, -0.5_347, -2.9_291, 9.1_435, -3.0_571, 8.9_764, 1.7_560, 8.7_358, -1.5_311]   )
							elif "trocr-large-handwritten" in checkpoint_url:
														__lowerCamelCase      :      Any         =torch.tensor(
														    [-2.6_437, -1.3_129, -2.2_596, -5.3_455, 6.3_539, 1.7_604, 5.4_991, 1.4_702, 5.6_113, 2.0_170]   )
							elif "trocr-base-printed" in checkpoint_url:
														__lowerCamelCase      :      Optional[Any]         =torch.tensor(
														    [-5.6_816, -5.8_388, 1.1_398, -6.9_034, 6.8_505, -2.4_393, 1.2_284, -1.0_232, -1.9_661, -3.9_210]   )
							elif "trocr-large-printed" in checkpoint_url:
														__lowerCamelCase      :      Optional[Any]         =torch.tensor(
														    [-6.0_162, -7.0_959, 4.4_155, -5.1_063, 7.0_468, -3.1_631, 2.6_466, -0.3_081, -0.8_106, -1.7_535]   )

							if "stage1" not in checkpoint_url:
														assert logits.shape == expected_shape, "Shape of logits not as expected"
														assert torch.allclose(logits[0, 0, :10]    ,						SCREAMING_SNAKE_CASE    ,						atol=1E-3   ), "First elements of logits not as expected"

							Path(SCREAMING_SNAKE_CASE   ).mkdir(exist_ok=SCREAMING_SNAKE_CASE   )
							print(F'Saving model to {pytorch_dump_folder_path}'   )
							model.save_pretrained(SCREAMING_SNAKE_CASE   )
							print(F'Saving processor to {pytorch_dump_folder_path}'   )
							processor.save_pretrained(SCREAMING_SNAKE_CASE   )


if __name__ == "__main__":
							_UpperCamelCase							=      argparse.ArgumentParser()

							parser.add_argument(
							    '--checkpoint_url',
							    default='https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt',
							    type=str,
							    help='URL to the original PyTorch checkpoint (.pth file).',
							)
							parser.add_argument(
							    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the folder to output PyTorch model.'
							)
							_UpperCamelCase							=      parser.parse_args()
							convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




