






from __future__ import annotations

import unittest

from transformers import 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 import (
			    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TFXLMForMultipleChoice,
			    TFXLMForQuestionAnsweringSimple,
			    TFXLMForSequenceClassification,
			    TFXLMForTokenClassification,
			    TFXLMModel,
			    TFXLMWithLMHeadModel,
			    XLMConfig,
			)


class _lowerCAmelCase	:







						def __init__(						self      ,							lowerCAmelCase_      ,							)		->       str:
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						parent
										_SCREAMING_SNAKE_CASE      :		Dict							=						1_3
										_SCREAMING_SNAKE_CASE      :		Any							=						7
										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						True
										_SCREAMING_SNAKE_CASE      :		str							=						True
										_SCREAMING_SNAKE_CASE      :		Dict							=						True
										_SCREAMING_SNAKE_CASE      :		List[str]							=						True
										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						True
										_SCREAMING_SNAKE_CASE      :		Tuple							=						False
										_SCREAMING_SNAKE_CASE      :		Tuple							=						False
										_SCREAMING_SNAKE_CASE      :		Dict							=						False
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						2
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						9_9
										_SCREAMING_SNAKE_CASE      :		int							=						0
										_SCREAMING_SNAKE_CASE      :		Tuple							=						3_2
										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						2
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						4
										_SCREAMING_SNAKE_CASE      :		str							=						0.1
										_SCREAMING_SNAKE_CASE      :		Any							=						0.1
										_SCREAMING_SNAKE_CASE      :		str							=						5_1_2
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						1_6
										_SCREAMING_SNAKE_CASE      :		int							=						2
										_SCREAMING_SNAKE_CASE      :		List[Any]							=						0.02
										_SCREAMING_SNAKE_CASE      :		Tuple							=						3
										_SCREAMING_SNAKE_CASE      :		int							=						4
										_SCREAMING_SNAKE_CASE      :		Any							=						'last'
										_SCREAMING_SNAKE_CASE      :		Any							=						True
										_SCREAMING_SNAKE_CASE      :		Any							=						None
										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						0







						def  A			(						self       )		->       Dict:
										_SCREAMING_SNAKE_CASE      :		Any							=						ids_tensor([self.batch_size, self.seq_length]      ,							self.vocab_size       )
										_SCREAMING_SNAKE_CASE      :		Tuple							=						random_attention_mask([self.batch_size, self.seq_length]      ,							dtype=tf.floataa       )

										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						None
										if self.use_input_lengths:
														_SCREAMING_SNAKE_CASE      :		Optional[int]							=						(
														    ids_tensor([self.batch_size]      ,							vocab_size=2       ) + self.seq_length - 2
														)  # small variation of seq_length

										_SCREAMING_SNAKE_CASE      :		Any							=						None
										if self.use_token_type_ids:
														_SCREAMING_SNAKE_CASE      :		Tuple							=						ids_tensor([self.batch_size, self.seq_length]      ,							self.n_langs       )

										_SCREAMING_SNAKE_CASE      :		Dict							=						None
										_SCREAMING_SNAKE_CASE      :		int							=						None
										_SCREAMING_SNAKE_CASE      :		Any							=						None
										if self.use_labels:
														_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						ids_tensor([self.batch_size]      ,							self.type_sequence_label_size       )
														_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						ids_tensor([self.batch_size, self.seq_length]      ,							self.num_labels       )
														_SCREAMING_SNAKE_CASE      :		List[Any]							=						ids_tensor([self.batch_size]      ,							2      ,							dtype=tf.floataa       )
														_SCREAMING_SNAKE_CASE      :		str							=						ids_tensor([self.batch_size]      ,							self.num_choices       )

										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						XLMConfig(
										    vocab_size=self.vocab_size      ,							n_special=self.n_special      ,							emb_dim=self.hidden_size      ,							n_layers=self.num_hidden_layers      ,							n_heads=self.num_attention_heads      ,							dropout=self.hidden_dropout_prob      ,							attention_dropout=self.attention_probs_dropout_prob      ,							gelu_activation=self.gelu_activation      ,							sinusoidal_embeddings=self.sinusoidal_embeddings      ,							asm=self.asm      ,							causal=self.causal      ,							n_langs=self.n_langs      ,							max_position_embeddings=self.max_position_embeddings      ,							initializer_range=self.initializer_range      ,							summary_type=self.summary_type      ,							use_proj=self.use_proj      ,							bos_token_id=self.bos_token_id      ,							)

										return (
										    config,
										    input_ids,
										    token_type_ids,
										    input_lengths,
										    sequence_labels,
										    token_labels,
										    is_impossible_labels,
										    choice_labels,
										    input_mask,
										)







						def  A			(						self      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							)		->       List[str]:
										_SCREAMING_SNAKE_CASE      :		int							=						TFXLMModel(config=lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		str							=						{'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
										_SCREAMING_SNAKE_CASE      :		Tuple							=						model(lowerCAmelCase_       )

										_SCREAMING_SNAKE_CASE      :		Any							=						[input_ids, input_mask]
										_SCREAMING_SNAKE_CASE      :		List[str]							=						model(lowerCAmelCase_       )
										self.parent.assertEqual(result.last_hidden_state.shape      ,							(self.batch_size, self.seq_length, self.hidden_size)       )







						def  A			(						self      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							)		->       str:
										_SCREAMING_SNAKE_CASE      :		Dict							=						TFXLMWithLMHeadModel(lowerCAmelCase_       )

										_SCREAMING_SNAKE_CASE      :		Tuple							=						{'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
										_SCREAMING_SNAKE_CASE      :		Dict							=						model(lowerCAmelCase_       )

										_SCREAMING_SNAKE_CASE      :		str							=						outputs

										self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.seq_length, self.vocab_size)       )







						def  A			(						self      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							)		->       Union[str, Any]:
										_SCREAMING_SNAKE_CASE      :		Any							=						TFXLMForQuestionAnsweringSimple(lowerCAmelCase_       )

										_SCREAMING_SNAKE_CASE      :		str							=						{'input_ids': input_ids, 'lengths': input_lengths}

										_SCREAMING_SNAKE_CASE      :		Tuple							=						model(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  A			(						self      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							)		->       Tuple:
										_SCREAMING_SNAKE_CASE      :		Dict							=						TFXLMForSequenceClassification(lowerCAmelCase_       )

										_SCREAMING_SNAKE_CASE      :		int							=						{'input_ids': input_ids, 'lengths': input_lengths}

										_SCREAMING_SNAKE_CASE      :		Tuple							=						model(lowerCAmelCase_       )

										self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.type_sequence_label_size)       )







						def  A			(						self      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							)		->       Optional[Any]:
										_SCREAMING_SNAKE_CASE      :		Any							=						self.num_labels
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						TFXLMForTokenClassification(config=lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		List[Any]							=						{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
										_SCREAMING_SNAKE_CASE      :		Dict							=						model(lowerCAmelCase_       )
										self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.seq_length, self.num_labels)       )







						def  A			(						self      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							)		->       Dict:
										_SCREAMING_SNAKE_CASE      :		Tuple							=						self.num_choices
										_SCREAMING_SNAKE_CASE      :		Any							=						TFXLMForMultipleChoice(config=lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						tf.tile(tf.expand_dims(lowerCAmelCase_      ,							1       )      ,							(1, self.num_choices, 1)       )
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						tf.tile(tf.expand_dims(lowerCAmelCase_      ,							1       )      ,							(1, self.num_choices, 1)       )
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						tf.tile(tf.expand_dims(lowerCAmelCase_      ,							1       )      ,							(1, self.num_choices, 1)       )
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						{
										    'input_ids': multiple_choice_inputs_ids,
										    'attention_mask': multiple_choice_input_mask,
										    'token_type_ids': multiple_choice_token_type_ids,
										}
										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						model(lowerCAmelCase_       )
										self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.num_choices)       )







						def  A			(						self       )		->       Union[str, Any]:
										_SCREAMING_SNAKE_CASE      :		str							=						self.prepare_config_and_inputs()
										(
										    _SCREAMING_SNAKE_CASE	
										)      :		Optional[Any]							=						config_and_inputs
										_SCREAMING_SNAKE_CASE      :		int							=						{
										    'input_ids': input_ids,
										    'token_type_ids': token_type_ids,
										    'langs': token_type_ids,
										    'lengths': input_lengths,
										}
										return config, inputs_dict







@require_tf
class _lowerCAmelCase	(						__UpperCAmelCase      ,     __UpperCAmelCase      ,     unittest.TestCase					):
						SCREAMING_SNAKE_CASE_:				Optional[Any]   				=							(
						    (
						        TFXLMModel,
						        TFXLMWithLMHeadModel,
						        TFXLMForSequenceClassification,
						        TFXLMForQuestionAnsweringSimple,
						        TFXLMForTokenClassification,
						        TFXLMForMultipleChoice,
						    )
						    if is_tf_available()
						    else ()
						)
						SCREAMING_SNAKE_CASE_:				Dict   				=							(
						    (TFXLMWithLMHeadModel,) if is_tf_available() else ()
						)  # TODO (PVP): Check other models whether language generation is also applicable
						SCREAMING_SNAKE_CASE_:				Any   				=							(
						    {
						        'feature-extraction': TFXLMModel,
						        'fill-mask': TFXLMWithLMHeadModel,
						        'question-answering': TFXLMForQuestionAnsweringSimple,
						        'text-classification': TFXLMForSequenceClassification,
						        'text-generation': TFXLMWithLMHeadModel,
						        'token-classification': TFXLMForTokenClassification,
						        'zero-shot': TFXLMForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						SCREAMING_SNAKE_CASE_:				Optional[Any]   				=							False
						SCREAMING_SNAKE_CASE_:				Tuple   				=							False







						def  A			(						self      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_      ,							lowerCAmelCase_       )		->       List[str]:
										if (
										    pipeline_test_casse_name == "QAPipelineTests"
										    and tokenizer_name is not None
										    and not tokenizer_name.endswith('Fast'       )
										):
														# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
														# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
														# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
														return True

										return False







						def  A			(						self       )		->       Dict:
										_SCREAMING_SNAKE_CASE      :		Dict							=						TFXLMModelTester(self       )
										_SCREAMING_SNAKE_CASE      :		List[str]							=						ConfigTester(self      ,							config_class=lowerCAmelCase_      ,							emb_dim=3_7       )







						def  A			(						self       )		->       Tuple:
										self.config_tester.run_common_tests()







						def  A			(						self       )		->       Any:
										_SCREAMING_SNAKE_CASE      :		Dict							=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_model(*lowerCAmelCase_       )







						def  A			(						self       )		->       List[Any]:
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_lm_head(*lowerCAmelCase_       )







						def  A			(						self       )		->       Union[str, Any]:
										_SCREAMING_SNAKE_CASE      :		Any							=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_qa(*lowerCAmelCase_       )







						def  A			(						self       )		->       Any:
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_sequence_classif(*lowerCAmelCase_       )







						def  A			(						self       )		->       Optional[Any]:
										_SCREAMING_SNAKE_CASE      :		Dict							=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_for_token_classification(*lowerCAmelCase_       )







						def  A			(						self       )		->       Union[str, Any]:
										_SCREAMING_SNAKE_CASE      :		int							=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_for_multiple_choice(*lowerCAmelCase_       )







						@slow
						def  A			(						self       )		->       Dict:
										for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_SCREAMING_SNAKE_CASE      :		str							=						TFXLMModel.from_pretrained(lowerCAmelCase_       )
														self.assertIsNotNone(lowerCAmelCase_       )







@require_tf
class _lowerCAmelCase	(						unittest.TestCase					):







						@slow
						def  A			(						self       )		->       Union[str, Any]:
										_SCREAMING_SNAKE_CASE      :		List[str]							=						TFXLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048'       )
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						tf.convert_to_tensor([[1_4, 4_4_7]]      ,							dtype=tf.intaa       )  # the president
										_SCREAMING_SNAKE_CASE      :		str							=						[
										    1_4,
										    4_4_7,
										    1_4,
										    4_4_7,
										    1_4,
										    4_4_7,
										    1_4,
										    4_4_7,
										    1_4,
										    4_4_7,
										    1_4,
										    4_4_7,
										    1_4,
										    4_4_7,
										    1_4,
										    4_4_7,
										    1_4,
										    4_4_7,
										    1_4,
										    4_4_7,
										]  # the president the president the president the president the president the president the president the president the president the president
										# TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						model.generate(lowerCAmelCase_      ,							do_sample=lowerCAmelCase_       )
										self.assertListEqual(output_ids[0].numpy().tolist()      ,							lowerCAmelCase_       )









"""simple docstring"""
from .chunk_utils import chunk_layer
from .data_transforms import make_atomaa_masks
from .feats import atomaa_to_atomaa, frames_and_literature_positions_to_atomaa_pos, torsion_angles_to_frames
from .loss import compute_predicted_aligned_error, compute_tm
from .protein import Protein as OFProtein
from .protein import to_pdb
from .rigid_utils import Rigid, Rotation
from .tensor_utils import dict_multimap, flatten_final_dims, permute_final_dims

