"""simple docstring"""






from __future__ import annotations

import unittest

import numpy as np

from transformers import LayoutLMConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers.models.layoutlm.modeling_tf_layoutlm import (
						    TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFLayoutLMForMaskedLM,
						    TFLayoutLMForQuestionAnswering,
						    TFLayoutLMForSequenceClassification,
						    TFLayoutLMForTokenClassification,
						    TFLayoutLMModel,
						)
class        snake_case :




					"""simple docstring"""

					def __init__(       self     :			Optional[Any]					,lowerCamelCase__     :			str					,lowerCamelCase__     :			Optional[Any]=13					,lowerCamelCase__     :			Any=7					,lowerCamelCase__     :			int=True					,lowerCamelCase__     :			str=True					,lowerCamelCase__     :			Dict=True					,lowerCamelCase__     :			Optional[Any]=True					,lowerCamelCase__     :			Any=99					,lowerCamelCase__     :			List[Any]=32					,lowerCamelCase__     :			Optional[int]=2					,lowerCamelCase__     :			Optional[int]=4					,lowerCamelCase__     :			Optional[Any]=37					,lowerCamelCase__     :			str="gelu"					,lowerCamelCase__     :			List[Any]=0.1					,lowerCamelCase__     :			Tuple=0.1					,lowerCamelCase__     :			Optional[int]=512					,lowerCamelCase__     :			Dict=16					,lowerCamelCase__     :			Union[str, Any]=2					,lowerCamelCase__     :			Optional[int]=0.0_2					,lowerCamelCase__     :			str=3					,lowerCamelCase__     :			Optional[Any]=4					,lowerCamelCase__     :			Any=None					,lowerCamelCase__     :			Union[str, Any]=1_000					,):
									UpperCAmelCase__    =							parent
									UpperCAmelCase__    =							batch_size
									UpperCAmelCase__    =							seq_length
									UpperCAmelCase__    =							is_training
									UpperCAmelCase__    =							use_input_mask
									UpperCAmelCase__    =							use_token_type_ids
									UpperCAmelCase__    =							use_labels
									UpperCAmelCase__    =							vocab_size
									UpperCAmelCase__    =							hidden_size
									UpperCAmelCase__    =							num_hidden_layers
									UpperCAmelCase__    =							num_attention_heads
									UpperCAmelCase__    =							intermediate_size
									UpperCAmelCase__    =							hidden_act
									UpperCAmelCase__    =							hidden_dropout_prob
									UpperCAmelCase__    =							attention_probs_dropout_prob
									UpperCAmelCase__    =							max_position_embeddings
									UpperCAmelCase__    =							type_vocab_size
									UpperCAmelCase__    =							type_sequence_label_size
									UpperCAmelCase__    =							initializer_range
									UpperCAmelCase__    =							num_labels
									UpperCAmelCase__    =							num_choices
									UpperCAmelCase__    =							scope
									UpperCAmelCase__    =							range_bbox







					def 	__lowerCAmelCase		(       self     :			str    ):
									UpperCAmelCase__    =							ids_tensor([self.batch_size, self.seq_length]					,self.vocab_size    )

									# convert bbox to numpy since TF does not support item assignment
									UpperCAmelCase__    =							ids_tensor([self.batch_size, self.seq_length, 4]					,self.range_bbox    ).numpy()
									# Ensure that bbox is legal
									for i in range(bbox.shape[0]    ):
													for j in range(bbox.shape[1]    ):
																	if bbox[i, j, 3] < bbox[i, j, 1]:
																					UpperCAmelCase__    =							bbox[i, j, 3]
																					UpperCAmelCase__    =							bbox[i, j, 1]
																					UpperCAmelCase__    =							t
																	if bbox[i, j, 2] < bbox[i, j, 0]:
																					UpperCAmelCase__    =							bbox[i, j, 2]
																					UpperCAmelCase__    =							bbox[i, j, 0]
																					UpperCAmelCase__    =							t
									UpperCAmelCase__    =							tf.convert_to_tensor(lowerCamelCase__    )

									UpperCAmelCase__    =							None
									if self.use_input_mask:
													UpperCAmelCase__    =							random_attention_mask([self.batch_size, self.seq_length]    )

									UpperCAmelCase__    =							None
									if self.use_token_type_ids:
													UpperCAmelCase__    =							ids_tensor([self.batch_size, self.seq_length]					,self.type_vocab_size    )

									UpperCAmelCase__    =							None
									UpperCAmelCase__    =							None
									UpperCAmelCase__    =							None
									if self.use_labels:
													UpperCAmelCase__    =							ids_tensor([self.batch_size]					,self.type_sequence_label_size    )
													UpperCAmelCase__    =							ids_tensor([self.batch_size, self.seq_length]					,self.num_labels    )
													UpperCAmelCase__    =							ids_tensor([self.batch_size]					,self.num_choices    )

									UpperCAmelCase__    =							LayoutLMConfig(
									    vocab_size=self.vocab_size					,hidden_size=self.hidden_size					,num_hidden_layers=self.num_hidden_layers					,num_attention_heads=self.num_attention_heads					,intermediate_size=self.intermediate_size					,hidden_act=self.hidden_act					,hidden_dropout_prob=self.hidden_dropout_prob					,attention_probs_dropout_prob=self.attention_probs_dropout_prob					,max_position_embeddings=self.max_position_embeddings					,type_vocab_size=self.type_vocab_size					,initializer_range=self.initializer_range					,)

									return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







					def 	__lowerCAmelCase		(       self     :			Union[str, Any]					,lowerCamelCase__     :			int					,lowerCamelCase__     :			Tuple					,lowerCamelCase__     :			List[Any]					,lowerCamelCase__     :			Dict					,lowerCamelCase__     :			List[str]					,lowerCamelCase__     :			List[str]					,lowerCamelCase__     :			Optional[Any]					,lowerCamelCase__     :			Optional[int]    ):
									UpperCAmelCase__    =							TFLayoutLMModel(config=lowerCamelCase__    )

									UpperCAmelCase__    =							model(lowerCamelCase__					,lowerCamelCase__					,attention_mask=lowerCamelCase__					,token_type_ids=lowerCamelCase__    )
									UpperCAmelCase__    =							model(lowerCamelCase__					,lowerCamelCase__					,token_type_ids=lowerCamelCase__    )
									UpperCAmelCase__    =							model(lowerCamelCase__					,lowerCamelCase__    )
									self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, self.seq_length, self.hidden_size)    )
									self.parent.assertEqual(result.pooler_output.shape					,(self.batch_size, self.hidden_size)    )







					def 	__lowerCAmelCase		(       self     :			str					,lowerCamelCase__     :			Any					,lowerCamelCase__     :			Optional[int]					,lowerCamelCase__     :			Tuple					,lowerCamelCase__     :			Optional[int]					,lowerCamelCase__     :			List[str]					,lowerCamelCase__     :			Any					,lowerCamelCase__     :			Optional[int]					,lowerCamelCase__     :			Dict    ):
									UpperCAmelCase__    =							TFLayoutLMForMaskedLM(config=lowerCamelCase__    )

									UpperCAmelCase__    =							model(lowerCamelCase__					,lowerCamelCase__					,attention_mask=lowerCamelCase__					,token_type_ids=lowerCamelCase__					,labels=lowerCamelCase__    )
									self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.seq_length, self.vocab_size)    )







					def 	__lowerCAmelCase		(       self     :			str					,lowerCamelCase__     :			Union[str, Any]					,lowerCamelCase__     :			Optional[int]					,lowerCamelCase__     :			int					,lowerCamelCase__     :			List[str]					,lowerCamelCase__     :			List[str]					,lowerCamelCase__     :			Optional[Any]					,lowerCamelCase__     :			List[Any]					,lowerCamelCase__     :			Optional[int]    ):
									UpperCAmelCase__    =							self.num_labels
									UpperCAmelCase__    =							TFLayoutLMForSequenceClassification(config=lowerCamelCase__    )

									UpperCAmelCase__    =							model(lowerCamelCase__					,lowerCamelCase__					,attention_mask=lowerCamelCase__					,token_type_ids=lowerCamelCase__    )
									self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.num_labels)    )







					def 	__lowerCAmelCase		(       self     :			Optional[int]					,lowerCamelCase__     :			int					,lowerCamelCase__     :			Optional[int]					,lowerCamelCase__     :			Tuple					,lowerCamelCase__     :			Optional[int]					,lowerCamelCase__     :			List[str]					,lowerCamelCase__     :			int					,lowerCamelCase__     :			Optional[int]					,lowerCamelCase__     :			Union[str, Any]    ):
									UpperCAmelCase__    =							self.num_labels
									UpperCAmelCase__    =							TFLayoutLMForTokenClassification(config=lowerCamelCase__    )

									UpperCAmelCase__    =							model(lowerCamelCase__					,lowerCamelCase__					,attention_mask=lowerCamelCase__					,token_type_ids=lowerCamelCase__					,labels=lowerCamelCase__    )
									self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.seq_length, self.num_labels)    )







					def 	__lowerCAmelCase		(       self     :			int					,lowerCamelCase__     :			List[str]					,lowerCamelCase__     :			Optional[Any]					,lowerCamelCase__     :			int					,lowerCamelCase__     :			str					,lowerCamelCase__     :			Tuple					,lowerCamelCase__     :			Optional[int]					,lowerCamelCase__     :			List[Any]					,lowerCamelCase__     :			int    ):
									UpperCAmelCase__    =							TFLayoutLMForQuestionAnswering(config=lowerCamelCase__    )

									UpperCAmelCase__    =							model(lowerCamelCase__					,lowerCamelCase__					,attention_mask=lowerCamelCase__					,token_type_ids=lowerCamelCase__    )
									self.parent.assertEqual(result.start_logits.shape					,(self.batch_size, self.seq_length)    )
									self.parent.assertEqual(result.end_logits.shape					,(self.batch_size, self.seq_length)    )







					def 	__lowerCAmelCase		(       self     :			Union[str, Any]    ):
									UpperCAmelCase__    =							self.prepare_config_and_inputs()
									(
									    (
									    UpperCAmelCase__
									)   ,			(
									    UpperCAmelCase__
									)   ,			(
									    UpperCAmelCase__
									)   ,			(
									    UpperCAmelCase__
									)   ,			(
									    UpperCAmelCase__
									)   ,			(
									    UpperCAmelCase__
									)   ,			(
									    UpperCAmelCase__
									)   ,			(
									    UpperCAmelCase__
									)   ,			
									)    =							config_and_inputs
									UpperCAmelCase__    =							{
									    'input_ids': input_ids,
									    'bbox': bbox,
									    'token_type_ids': token_type_ids,
									    'attention_mask': input_mask,
									}
									return config, inputs_dict






@require_tf
class        snake_case (				__UpperCAmelCase   ,					__UpperCAmelCase   ,					unittest.TestCase       ):




					"""simple docstring"""

					snake_case__       					=				(
					    (
					        TFLayoutLMModel,
					        TFLayoutLMForMaskedLM,
					        TFLayoutLMForTokenClassification,
					        TFLayoutLMForSequenceClassification,
					        TFLayoutLMForQuestionAnswering,
					    )
					    if is_tf_available()
					    else ()
					)
					snake_case__       					=				(
					    {
					        "feature-extraction": TFLayoutLMModel,
					        "fill-mask": TFLayoutLMForMaskedLM,
					        "text-classification": TFLayoutLMForSequenceClassification,
					        "token-classification": TFLayoutLMForTokenClassification,
					        "zero-shot": TFLayoutLMForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					snake_case__       					=				False
					snake_case__       					=				True
					snake_case__       					=				10







					def 	__lowerCAmelCase		(       self     :			List[Any]    ):
									UpperCAmelCase__    =							TFLayoutLMModelTester(self    )
									UpperCAmelCase__    =							ConfigTester(self					,config_class=lowerCamelCase__					,hidden_size=37    )







					def 	__lowerCAmelCase		(       self     :			Optional[Any]    ):
									self.config_tester.run_common_tests()







					def 	__lowerCAmelCase		(       self     :			List[Any]    ):
									UpperCAmelCase__    =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowerCamelCase__    )







					def 	__lowerCAmelCase		(       self     :			Optional[int]    ):
									UpperCAmelCase__    =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase__    )







					def 	__lowerCAmelCase		(       self     :			Dict    ):
									UpperCAmelCase__    =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase__    )







					def 	__lowerCAmelCase		(       self     :			Dict    ):
									UpperCAmelCase__    =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*lowerCamelCase__    )







					def 	__lowerCAmelCase		(       self     :			Optional[int]    ):
									UpperCAmelCase__    =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*lowerCamelCase__    )







					@slow
					def 	__lowerCAmelCase		(       self     :			Tuple    ):
									for model_name in TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													UpperCAmelCase__    =							TFLayoutLMModel.from_pretrained(lowerCamelCase__    )
													self.assertIsNotNone(lowerCamelCase__    )







					@unittest.skip('Onnx compliancy broke with TF 2.10'    )
					def 	__lowerCAmelCase		(       self     :			Optional[Any]    ):
									pass






def 						a_		(      ):
				# Here we prepare a batch of 2 sequences to test a LayoutLM forward pass on:
				# fmt: off
				UpperCAmelCase__    =							tf.convert_to_tensor([[1_0_1,1_0_1_9,1_0_1_4,1_0_1_6,1_0_3_7,1_2_8_4_9,4_7_4_7,1_0_0_4,1_4_2_4_6,2_2_7_8,5_4_3_9,4_5_2_4,5_0_0_2,2_9_3_0,2_1_9_3,2_9_3_0,4_3_4_1,3_2_0_8,1_0_0_5,1_0_5_5,2_1_7_1,2_8_4_8,1_1_3_0_0,3_5_3_1,1_0_2],[1_0_1,4_0_7_0,4_0_3_4,7_0_2_0,1_0_2_4,3_0_5_8,1_0_1_5,1_0_1_3,2_8_6_1,1_0_1_3,6_0_7_0,1_9_2_7_4,2_7_7_2,6_2_0_5,2_7_8_1_4,1_6_1_4_7,1_6_1_4_7,4_3_4_3,2_0_4_7,1_0_2_8_3,1_0_9_6_9,1_4_3_8_9,1_0_1_2,2_3_3_8,1_0_2]] )  # noqa: E231
				UpperCAmelCase__    =							tf.convert_to_tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],] )  # noqa: E231
				UpperCAmelCase__    =							tf.convert_to_tensor([[[0,0,0,0],[4_2_3,2_3_7,4_4_0,2_5_1],[4_2_7,2_7_2,4_4_1,2_8_7],[4_1_9,1_1_5,4_3_7,1_2_9],[9_6_1,8_8_5,9_9_2,9_1_2],[2_5_6,3_8,3_3_0,5_8],[2_5_6,3_8,3_3_0,5_8],[3_3_6,4_2,3_5_3,5_7],[3_6_0,3_9,4_0_1,5_6],[3_6_0,3_9,4_0_1,5_6],[4_1_1,3_9,4_7_1,5_9],[4_7_9,4_1,5_2_8,5_9],[5_3_3,3_9,6_3_0,6_0],[6_7,1_1_3,1_3_4,1_3_1],[1_4_1,1_1_5,2_0_9,1_3_2],[6_8,1_4_9,1_3_3,1_6_6],[1_4_1,1_4_9,1_8_7,1_6_4],[1_9_5,1_4_8,2_8_7,1_6_5],[1_9_5,1_4_8,2_8_7,1_6_5],[1_9_5,1_4_8,2_8_7,1_6_5],[2_9_5,1_4_8,3_4_9,1_6_5],[4_4_1,1_4_9,4_9_2,1_6_6],[4_9_7,1_4_9,5_4_6,1_6_4],[6_4,2_0_1,1_2_5,2_1_8],[1_0_0_0,1_0_0_0,1_0_0_0,1_0_0_0]],[[0,0,0,0],[6_6_2,1_5_0,7_5_4,1_6_6],[6_6_5,1_9_9,7_4_2,2_1_1],[5_1_9,2_1_3,5_5_4,2_2_8],[5_1_9,2_1_3,5_5_4,2_2_8],[1_3_4,4_3_3,1_8_7,4_5_4],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[3_1_4,4_6_9,3_7_6,4_8_2],[5_0_4,6_8_4,5_8_2,7_0_6],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[6_1_0,7_4_9,6_5_2,7_6_5],[1_3_0,6_5_9,1_6_8,6_7_2],[1_7_6,6_5_7,2_3_7,6_7_2],[2_3_8,6_5_7,3_1_2,6_7_2],[4_4_3,6_5_3,6_2_8,6_7_2],[4_4_3,6_5_3,6_2_8,6_7_2],[7_1_6,3_0_1,8_2_5,3_1_7],[1_0_0_0,1_0_0_0,1_0_0_0,1_0_0_0]]] )  # noqa: E231
				UpperCAmelCase__    =							tf.convert_to_tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]] )  # noqa: E231
				# these are sequence labels (i.e. at the token level)
				UpperCAmelCase__    =							tf.convert_to_tensor([[-1_0_0,1_0,1_0,1_0,9,1,-1_0_0,7,7,-1_0_0,7,7,4,2,5,2,8,8,-1_0_0,-1_0_0,5,0,3,2,-1_0_0],[-1_0_0,1_2,1_2,1_2,-1_0_0,1_2,1_0,-1_0_0,-1_0_0,-1_0_0,-1_0_0,1_0,1_2,9,-1_0_0,-1_0_0,-1_0_0,1_0,1_0,1_0,9,1_2,-1_0_0,1_0,-1_0_0]] )  # noqa: E231
				# fmt: on

				return input_ids, attention_mask, bbox, token_type_ids, labels







@require_tf
class        snake_case (				unittest.TestCase       ):




					"""simple docstring"""

					@slow
					def 	__lowerCAmelCase		(       self     :			List[Any]    ):
									UpperCAmelCase__    =							TFLayoutLMModel.from_pretrained('microsoft/layoutlm-base-uncased'    )

									UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__    =							prepare_layoutlm_batch_inputs()

									# forward pass
									UpperCAmelCase__    =							model(input_ids=lowerCamelCase__					,bbox=lowerCamelCase__					,attention_mask=lowerCamelCase__					,token_type_ids=lowerCamelCase__    )

									# test the sequence output on [0, :3, :3]
									UpperCAmelCase__    =							tf.convert_to_tensor(
									    [[0.1_7_8_5, -0.1_9_4_7, -0.0_4_2_5], [-0.3_2_5_4, -0.2_8_0_7, 0.2_5_5_3], [-0.5_3_9_1, -0.3_3_2_2, 0.3_3_6_4]]					,)

									self.assertTrue(np.allclose(outputs.last_hidden_state[0, :3, :3]					,lowerCamelCase__					,atol=1e-3    )    )

									# test the pooled output on [1, :3]
									UpperCAmelCase__    =							tf.convert_to_tensor([-0.6_5_8_0, -0.0_2_1_4, 0.8_5_5_2]    )

									self.assertTrue(np.allclose(outputs.pooler_output[1, :3]					,lowerCamelCase__					,atol=1e-3    )    )







					@slow
					def 	__lowerCAmelCase		(       self     :			Optional[Any]    ):
									# initialize model with randomly initialized sequence classification head
									UpperCAmelCase__    =							TFLayoutLMForSequenceClassification.from_pretrained('microsoft/layoutlm-base-uncased'					,num_labels=2    )

									UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__    =							prepare_layoutlm_batch_inputs()

									# forward pass
									UpperCAmelCase__    =							model(
									    input_ids=lowerCamelCase__					,bbox=lowerCamelCase__					,attention_mask=lowerCamelCase__					,token_type_ids=lowerCamelCase__					,labels=tf.convert_to_tensor([1, 1]    )					,)

									# test whether we get a loss as a scalar
									UpperCAmelCase__    =							outputs.loss
									UpperCAmelCase__    =							(2,)
									self.assertEqual(loss.shape					,lowerCamelCase__    )

									# test the shape of the logits
									UpperCAmelCase__    =							outputs.logits
									UpperCAmelCase__    =							(2, 2)
									self.assertEqual(logits.shape					,lowerCamelCase__    )







					@slow
					def 	__lowerCAmelCase		(       self     :			int    ):
									# initialize model with randomly initialized token classification head
									UpperCAmelCase__    =							TFLayoutLMForTokenClassification.from_pretrained('microsoft/layoutlm-base-uncased'					,num_labels=13    )

									UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__    =							prepare_layoutlm_batch_inputs()

									# forward pass
									UpperCAmelCase__    =							model(
									    input_ids=lowerCamelCase__					,bbox=lowerCamelCase__					,attention_mask=lowerCamelCase__					,token_type_ids=lowerCamelCase__					,labels=lowerCamelCase__    )

									# test the shape of the logits
									UpperCAmelCase__    =							outputs.logits
									UpperCAmelCase__    =							tf.convert_to_tensor((2, 25, 13)    )
									self.assertEqual(logits.shape					,lowerCamelCase__    )







					@slow
					def 	__lowerCAmelCase		(       self     :			int    ):
									# initialize model with randomly initialized token classification head
									UpperCAmelCase__    =							TFLayoutLMForQuestionAnswering.from_pretrained('microsoft/layoutlm-base-uncased'    )

									UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__   ,			UpperCAmelCase__    =							prepare_layoutlm_batch_inputs()

									# forward pass
									UpperCAmelCase__    =							model(input_ids=lowerCamelCase__					,bbox=lowerCamelCase__					,attention_mask=lowerCamelCase__					,token_type_ids=lowerCamelCase__    )

									# test the shape of the logits
									UpperCAmelCase__    =							tf.convert_to_tensor((2, 25)    )
									self.assertEqual(outputs.start_logits.shape					,lowerCamelCase__    )
									self.assertEqual(outputs.end_logits.shape					,lowerCamelCase__    )




"""simple docstring"""






import argparse
from pathlib import Path
from typing import Dict, OrderedDict, Tuple

import torch
from audiocraft.models import MusicGen

from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
    EncodecModel,
    MusicgenDecoderConfig,
    MusicgenForConditionalGeneration,
    MusicgenProcessor,
    TaEncoderModel,
)
from transformers.models.musicgen.modeling_musicgen import MusicgenForCausalLM
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase__						:						Optional[int]		  =			logging.get_logger(__name__)


lowerCAmelCase__						:						Optional[Any]		  =			['model.decoder.embed_positions.weights']



def 						a_		(      lowerCamelCase ):
				if "emb" in name:
								UpperCAmelCase__    =							name.replace('emb'      ,    'model.decoder.embed_tokens' )
				if "transformer" in name:
								UpperCAmelCase__    =							name.replace('transformer'      ,    'model.decoder' )
				if "cross_attention" in name:
								UpperCAmelCase__    =							name.replace('cross_attention'      ,    'encoder_attn' )
				if "linear1" in name:
								UpperCAmelCase__    =							name.replace('linear1'      ,    'fc1' )
				if "linear2" in name:
								UpperCAmelCase__    =							name.replace('linear2'      ,    'fc2' )
				if "norm1" in name:
								UpperCAmelCase__    =							name.replace('norm1'      ,    'self_attn_layer_norm' )
				if "norm_cross" in name:
								UpperCAmelCase__    =							name.replace('norm_cross'      ,    'encoder_attn_layer_norm' )
				if "norm2" in name:
								UpperCAmelCase__    =							name.replace('norm2'      ,    'final_layer_norm' )
				if "out_norm" in name:
								UpperCAmelCase__    =							name.replace('out_norm'      ,    'model.decoder.layer_norm' )
				if "linears" in name:
								UpperCAmelCase__    =							name.replace('linears'      ,    'lm_heads' )
				if "condition_provider.conditioners.description.output_proj" in name:
								UpperCAmelCase__    =							name.replace('condition_provider.conditioners.description.output_proj'      ,    'enc_to_dec_proj' )
				return name



def 						a_		(      lowerCamelCase      ,    lowerCamelCase ):
				UpperCAmelCase__    =							list(state_dict.keys() )
				UpperCAmelCase__    =							{}
				for key in keys:
								UpperCAmelCase__    =							state_dict.pop(lowerCamelCase )
								UpperCAmelCase__    =							rename_keys(lowerCamelCase )
								if "in_proj_weight" in key:
												# split fused qkv proj
												UpperCAmelCase__    =							val[:hidden_size, :]
												UpperCAmelCase__    =							val[hidden_size : 2 * hidden_size, :]
												UpperCAmelCase__    =							val[-hidden_size:, :]
								elif "enc_to_dec_proj" in key:
												UpperCAmelCase__    =							val
								else:
												UpperCAmelCase__    =							val
				return state_dict, enc_dec_proj_state_dict



def 						a_		(      lowerCamelCase ):
				if checkpoint == "small":
								# default config values
								UpperCAmelCase__    =							1_0_2_4
								UpperCAmelCase__    =							2_4
								UpperCAmelCase__    =							1_6
				elif checkpoint == "medium":
								UpperCAmelCase__    =							1_5_3_6
								UpperCAmelCase__    =							4_8
								UpperCAmelCase__    =							2_4
				elif checkpoint == "large":
								UpperCAmelCase__    =							2_0_4_8
								UpperCAmelCase__    =							4_8
								UpperCAmelCase__    =							3_2
				else:
								raise ValueError(f'''Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.''' )
				UpperCAmelCase__    =							MusicgenDecoderConfig(
				    hidden_size=lowerCamelCase      ,    ffn_dim=hidden_size * 4      ,    num_hidden_layers=lowerCamelCase      ,    num_attention_heads=lowerCamelCase      ,    )
				return config



@torch.no_grad()
def 						a_		(      lowerCamelCase      ,    lowerCamelCase=None      ,    lowerCamelCase=None      ,    lowerCamelCase="cpu" ):
				UpperCAmelCase__    =							MusicGen.get_pretrained(lowerCamelCase      ,    device=lowerCamelCase )
				UpperCAmelCase__    =							decoder_config_from_checkpoint(lowerCamelCase )

				UpperCAmelCase__    =							fairseq_model.lm.state_dict()
				UpperCAmelCase__   ,			UpperCAmelCase__    =							rename_state_dict(
				    lowerCamelCase      ,    hidden_size=decoder_config.hidden_size )

				UpperCAmelCase__    =							TaEncoderModel.from_pretrained('t5-base' )
				UpperCAmelCase__    =							EncodecModel.from_pretrained('facebook/encodec_32khz' )
				UpperCAmelCase__    =							MusicgenForCausalLM(lowerCamelCase ).eval()

				# load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
				UpperCAmelCase__   ,			UpperCAmelCase__    =							decoder.load_state_dict(lowerCamelCase      ,    strict=lowerCamelCase )

				for key in missing_keys.copy():
								if key.startswith(('text_encoder', 'audio_encoder') ) or key in EXPECTED_MISSING_KEYS:
												missing_keys.remove(lowerCamelCase )

				if len(lowerCamelCase ) > 0:
								raise ValueError(f'''Missing key(s) in state_dict: {missing_keys}''' )

				if len(lowerCamelCase ) > 0:
								raise ValueError(f'''Unexpected key(s) in state_dict: {unexpected_keys}''' )

				# init the composite model
				UpperCAmelCase__    =							MusicgenForConditionalGeneration(text_encoder=lowerCamelCase      ,    audio_encoder=lowerCamelCase      ,    decoder=lowerCamelCase )

				# load the pre-trained enc-dec projection (from the decoder state dict)
				model.enc_to_dec_proj.load_state_dict(lowerCamelCase )

				# check we can do a forward pass
				UpperCAmelCase__    =							torch.arange(0      ,    8      ,    dtype=torch.long ).reshape(2      ,    -1 )
				UpperCAmelCase__    =							input_ids.reshape(2 * 4      ,    -1 )

				with torch.no_grad():
								UpperCAmelCase__    =							model(input_ids=lowerCamelCase      ,    decoder_input_ids=lowerCamelCase ).logits

				if logits.shape != (8, 1, 2_0_4_8):
								raise ValueError('Incorrect shape for logits' )

				# now construct the processor
				UpperCAmelCase__    =							AutoTokenizer.from_pretrained('t5-base' )
				UpperCAmelCase__    =							AutoFeatureExtractor.from_pretrained('facebook/encodec_32khz'      ,    padding_side='left' )

				UpperCAmelCase__    =							MusicgenProcessor(feature_extractor=lowerCamelCase      ,    tokenizer=lowerCamelCase )

				# set the appropriate bos/pad token ids
				UpperCAmelCase__    =							2_0_4_8
				UpperCAmelCase__    =							2_0_4_8

				# set other default generation config params
				UpperCAmelCase__    =							int(3_0 * audio_encoder.config.frame_rate )
				UpperCAmelCase__    =							True
				UpperCAmelCase__    =							3.0

				if pytorch_dump_folder is not None:
								Path(lowerCamelCase ).mkdir(exist_ok=lowerCamelCase )
								logger.info(f'''Saving model {checkpoint} to {pytorch_dump_folder}''' )
								model.save_pretrained(lowerCamelCase )
								processor.save_pretrained(lowerCamelCase )

				if repo_id:
								logger.info(f'''Pushing model {checkpoint} to {repo_id}''' )
								model.push_to_hub(lowerCamelCase )
								processor.push_to_hub(lowerCamelCase )


if __name__ == "__main__":
						lowerCAmelCase__						:						Tuple		  =			argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--checkpoint',
						    default='small',
						    type=str,
						    help='Checkpoint size of the MusicGen model you\'d like to convert. Can be one of: `[\'small\', \'medium\', \'large\']`.',
						)
						parser.add_argument(
						    '--pytorch_dump_folder',
						    required=True,
						    default=None,
						    type=str,
						    help='Path to the output PyTorch model directory.',
						)
						parser.add_argument(
						    '--push_to_hub', default=None, type=str, help='Where to upload the converted model on the 🤗 hub.'
						)
						parser.add_argument(
						    '--device', default='cpu', type=str, help='Torch device to run the conversion, either cpu or cuda.'
						)

						lowerCAmelCase__						:						List[str]		  =			parser.parse_args()
						convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)



