
'''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()
lowercase					:							Tuple   =      logging.get_logger(__name__)







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			Tuple:
	_snake_case     		=	[]
	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        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			str:
	for i in range(encoder_config.num_hidden_layers			):
		# queries, keys and values (only weights, no biases)
		_snake_case     		=	state_dict.pop(F'encoder.deit.blocks.{i}.attn.qkv.weight'			)

		_snake_case     		=	in_proj_weight[
		    : encoder_config.hidden_size, :
		]
		_snake_case     		=	in_proj_weight[
		    encoder_config.hidden_size : encoder_config.hidden_size * 2, :
		]
		_snake_case     		=	in_proj_weight[
		    -encoder_config.hidden_size :, :
		]







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A    ,       __A			) ->			int:
	_snake_case     		=	dct.pop(__A			)
	_snake_case     		=	val







def        SCREAMING_SNAKE_CASE__			(     __A			) ->			Optional[int]:
	if "handwritten" in checkpoint_url:
		_snake_case     		=	'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:
		_snake_case     		=	'https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg'
	_snake_case     		=	Image.open(requests.get(__A    ,       stream=__A			).raw			).convert('RGB'			)
	return im







@torch.no_grad()
def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			List[Any]:
	_snake_case     		=	ViTConfig(image_size=384    ,       qkv_bias=__A			)
	_snake_case     		=	TrOCRConfig()

	# size of the architecture
	if "base" in checkpoint_url:
		_snake_case     		=	768
	elif "large" in checkpoint_url:
		# use ViT-large encoder
		_snake_case     		=	1_024
		_snake_case     		=	4_096
		_snake_case     		=	24
		_snake_case     		=	16
		_snake_case     		=	1_024
	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:
		_snake_case     		=	False
		_snake_case     		=	'relu'
		_snake_case     		=	1_024
		_snake_case     		=	True
		_snake_case     		=	False
		_snake_case     		=	False

	# load HuggingFace model
	_snake_case     		=	ViTModel(__A    ,       add_pooling_layer=__A			)
	_snake_case     		=	TrOCRForCausalLM(__A			)
	_snake_case     		=	VisionEncoderDecoderModel(encoder=__A    ,       decoder=__A			)
	model.eval()

	# load state_dict of original model, rename some keys
	_snake_case     		=	torch.hub.load_state_dict_from_url(__A    ,       map_location='cpu'    ,       check_hash=__A			)['model']

	_snake_case     		=	create_rename_keys(__A    ,       __A			)
	for src, dest in rename_keys:
		rename_key(__A    ,       __A    ,       __A			)
	read_in_q_k_v(__A    ,       __A			)

	# 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():
		_snake_case     		=	state_dict.pop(__A			)
		if key.startswith('decoder'			) and "output_projection" not in key:
			_snake_case     		=	val
		else:
			_snake_case     		=	val

    # load state dict
	model.load_state_dict(__A			)

	# Check outputs on an image
	_snake_case     		=	ViTImageProcessor(size=encoder_config.image_size			)
	_snake_case     		=	RobertaTokenizer.from_pretrained('roberta-large'			)
	_snake_case     		=	TrOCRProcessor(__A    ,       __A			)

	_snake_case     		=	processor(images=prepare_img(__A			)    ,       return_tensors='pt'			).pixel_values

	# verify logits
	_snake_case     		=	torch.tensor([[model.config.decoder.decoder_start_token_id]]			)
	_snake_case     		=	model(pixel_values=__A    ,       decoder_input_ids=__A			)
	_snake_case     		=	outputs.logits

	_snake_case     		=	torch.Size([1, 1, 50_265]			)
	if "trocr-base-handwritten" in checkpoint_url:
		_snake_case     		=	torch.tensor(
		    [-1.4_5_0_2, -4.6_6_8_3, -0.5_3_4_7, -2.9_2_9_1, 9.1_4_3_5, -3.0_5_7_1, 8.9_7_6_4, 1.7_5_6_0, 8.7_3_5_8, -1.5_3_1_1]			)
	elif "trocr-large-handwritten" in checkpoint_url:
		_snake_case     		=	torch.tensor(
		    [-2.6_4_3_7, -1.3_1_2_9, -2.2_5_9_6, -5.3_4_5_5, 6.3_5_3_9, 1.7_6_0_4, 5.4_9_9_1, 1.4_7_0_2, 5.6_1_1_3, 2.0_1_7_0]			)
	elif "trocr-base-printed" in checkpoint_url:
		_snake_case     		=	torch.tensor(
		    [-5.6_8_1_6, -5.8_3_8_8, 1.1_3_9_8, -6.9_0_3_4, 6.8_5_0_5, -2.4_3_9_3, 1.2_2_8_4, -1.0_2_3_2, -1.9_6_6_1, -3.9_2_1_0]			)
	elif "trocr-large-printed" in checkpoint_url:
		_snake_case     		=	torch.tensor(
		    [-6.0_1_6_2, -7.0_9_5_9, 4.4_1_5_5, -5.1_0_6_3, 7.0_4_6_8, -3.1_6_3_1, 2.6_4_6_6, -0.3_0_8_1, -0.8_1_0_6, -1.7_5_3_5]			)

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

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


if __name__ == "__main__":
					lowercase					:							int   =      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."
					)
					lowercase					:							List[Any]   =      parser.parse_args()
					convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)



'''simple docstring'''







from __future__ import annotations

import math







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			list:
	if len(__A			) != 2 or len(a[0]			) != 2 or len(__A			) != 2 or len(b[0]			) != 2:
		raise Exception('Matrices are not 2x2'			)
	_snake_case     		=	[
	    [a[0][0] * b[0][0] + a[0][1] * b[1][0], a[0][0] * b[0][1] + a[0][1] * b[1][1]],
	    [a[1][0] * b[0][0] + a[1][1] * b[1][0], a[1][0] * b[0][1] + a[1][1] * b[1][1]],
	]
	return new_matrix







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			Dict:
	return [
	    [matrix_a[row][col] + matrix_b[row][col] for col in range(len(matrix_a[row]			)			)]
	    for row in range(len(__A			)			)
	]







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			Dict:
	return [
	    [matrix_a[row][col] - matrix_b[row][col] for col in range(len(matrix_a[row]			)			)]
	    for row in range(len(__A			)			)
	]







def        SCREAMING_SNAKE_CASE__			(     __A			) ->			tuple[list, list, list, list]:
	if len(__A			) % 2 != 0 or len(a[0]			) % 2 != 0:
		raise Exception('Odd matrices are not supported!'			)

	_snake_case     		=	len(__A			)
	_snake_case     		=	matrix_length // 2

	_snake_case     		=	[[a[i][j] for j in range(__A    ,       __A			)] for i in range(__A			)]
	_snake_case     		=	[
	    [a[i][j] for j in range(__A    ,       __A			)] for i in range(__A    ,       __A			)
	]

	_snake_case     		=	[[a[i][j] for j in range(__A			)] for i in range(__A			)]
	_snake_case     		=	[[a[i][j] for j in range(__A			)] for i in range(__A    ,       __A			)]

	return top_left, top_right, bot_left, bot_right







def        SCREAMING_SNAKE_CASE__			(     __A			) ->			tuple[int, int]:
	return len(__A			), len(matrix[0]			)







def        SCREAMING_SNAKE_CASE__			(     __A			) ->			None:
	print('\n'.join(str(__A			) for line in matrix			)			)







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			list:
	if matrix_dimensions(__A			) == (2, 2):
		return default_matrix_multiplication(__A    ,       __A			)

	_snake_case		,				_snake_case		,				_snake_case		,				_snake_case     		=	split_matrix(__A			)
	_snake_case		,				_snake_case		,				_snake_case		,				_snake_case     		=	split_matrix(__A			)

	_snake_case     		=	actual_strassen(__A    ,       matrix_subtraction(__A    ,       __A			)			)
	_snake_case     		=	actual_strassen(matrix_addition(__A    ,       __A			)    ,       __A			)
	_snake_case     		=	actual_strassen(matrix_addition(__A    ,       __A			)    ,       __A			)
	_snake_case     		=	actual_strassen(__A    ,       matrix_subtraction(__A    ,       __A			)			)
	_snake_case     		=	actual_strassen(matrix_addition(__A    ,       __A			)    ,       matrix_addition(__A    ,       __A			)			)
	_snake_case     		=	actual_strassen(matrix_subtraction(__A    ,       __A			)    ,       matrix_addition(__A    ,       __A			)			)
	_snake_case     		=	actual_strassen(matrix_subtraction(__A    ,       __A			)    ,       matrix_addition(__A    ,       __A			)			)

	_snake_case     		=	matrix_addition(matrix_subtraction(matrix_addition(__A    ,       __A			)    ,       __A			)    ,       __A			)
	_snake_case     		=	matrix_addition(__A    ,       __A			)
	_snake_case     		=	matrix_addition(__A    ,       __A			)
	_snake_case     		=	matrix_subtraction(matrix_subtraction(matrix_addition(__A    ,       __A			)    ,       __A			)    ,       __A			)

	# construct the new matrix from our 4 quadrants
	_snake_case     		=	[]
	for i in range(len(__A			)			):
		new_matrix.append(top_left[i] + top_right[i]			)
	for i in range(len(__A			)			):
		new_matrix.append(bot_left[i] + bot_right[i]			)
	return new_matrix







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			list:
	if matrix_dimensions(__A			)[1] != matrix_dimensions(__A			)[0]:
		_snake_case     		=	(
		    'Unable to multiply these matrices, please check the dimensions.\n'
		    F'Matrix A: {matrixa}\n'
		    F'Matrix B: {matrixa}'
		)
		raise Exception(__A			)
	_snake_case     		=	matrix_dimensions(__A			)
	_snake_case     		=	matrix_dimensions(__A			)

	if dimensiona[0] == dimensiona[1] and dimensiona[0] == dimensiona[1]:
		return [matrixa, matrixa]

	_snake_case     		=	max(*__A    ,       *__A			)
	_snake_case     		=	int(math.pow(2    ,       math.ceil(math.loga(__A			)			)			)			)
	_snake_case     		=	matrixa
	_snake_case     		=	matrixa

	# Adding zeros to the matrices so that the arrays dimensions are the same and also
	# power of 2
	for i in range(0    ,       __A			):
		if i < dimensiona[0]:
			for _ in range(dimensiona[1]    ,       __A			):
				new_matrixa[i].append(0			)
		else:
			new_matrixa.append([0] * maxim			)
		if i < dimensiona[0]:
			for _ in range(dimensiona[1]    ,       __A			):
				new_matrixa[i].append(0			)
		else:
			new_matrixa.append([0] * maxim			)

	_snake_case     		=	actual_strassen(__A    ,       __A			)

	# Removing the additional zeros
	for i in range(0    ,       __A			):
		if i < dimensiona[0]:
			for _ in range(dimensiona[1]    ,       __A			):
				final_matrix[i].pop()
		else:
			final_matrix.pop()
	return final_matrix


if __name__ == "__main__":
					lowercase					:							Any   =      [
					    [2, 3, 4, 5],
					    [6, 4, 3, 1],
					    [2, 3, 6, 7],
					    [3, 1, 2, 4],
					    [2, 3, 4, 5],
					    [6, 4, 3, 1],
					    [2, 3, 6, 7],
					    [3, 1, 2, 4],
					    [2, 3, 4, 5],
					    [6, 2, 3, 1],
					]
					lowercase					:							List[str]   =      [[0, 2, 1, 1], [16, 2, 3, 3], [2, 2, 7, 7], [13, 11, 22, 4]]
					print(strassen(matrixa, matrixa))

