




'''simple docstring'''






import os
import pickle
import tempfile
import unittest
from typing import Callable, Optional

import numpy as np

from transformers import (
    BatchEncoding,
    BertTokenizer,
    BertTokenizerFast,
    PreTrainedTokenizer,
    PreTrainedTokenizerFast,
    TensorType,
    TokenSpan,
    is_tokenizers_available,
)
from transformers.models.gpta.tokenization_gpta import GPTaTokenizer
from transformers.testing_utils import CaptureStderr, require_flax, require_tf, require_tokenizers, require_torch, slow


if is_tokenizers_available():
				from tokenizers import Tokenizer
				from tokenizers.models import WordPiece



class    A						(					unittest.TestCase						):




				'''simple docstring'''





				def 			a_			(self					,    _UpperCAmelCase     )      ->			Optional[int]:
							__UpperCamelCase :     Any										=  list(tokenizer_class.max_model_input_sizes.keys()     )
							for model_name in sa_models[:1]:
										__UpperCamelCase :     int										=  tokenizer_class.from_pretrained(_UpperCAmelCase     )
										self.assertIsNotNone(_UpperCAmelCase     )
										self.assertIsInstance(_UpperCAmelCase					,    _UpperCAmelCase     )
										self.assertIsInstance(_UpperCAmelCase					,    _UpperCAmelCase     )

										for special_tok in tokenizer.all_special_tokens:
													self.assertIsInstance(_UpperCAmelCase					,    _UpperCAmelCase     )
													__UpperCamelCase :     Tuple										=  tokenizer.convert_tokens_to_ids(_UpperCAmelCase     )
													self.assertIsInstance(_UpperCAmelCase					,    _UpperCAmelCase     )
				def 			a_			(self					,    _UpperCAmelCase					,    _UpperCAmelCase = None     )      ->			Tuple:
							__UpperCamelCase :     Union[str, Any]										=  pickle.dumps(_UpperCAmelCase     )
							self.assertIsNotNone(_UpperCAmelCase     )

							__UpperCamelCase :     Tuple										=  pickle.loads(_UpperCAmelCase     )

							# Ensure is_fast is correctly restored
							self.assertEqual(be_restored.is_fast					,    be_original.is_fast     )

							# Ensure encodings are potentially correctly restored
							if be_original.is_fast:
										self.assertIsNotNone(be_restored.encodings     )
							else:
										self.assertIsNone(be_restored.encodings     )

							# Ensure the keys are the same
							for original_v, restored_v in zip(be_original.values()					,    be_restored.values()     ):
										if equal_op:
													self.assertTrue(equal_op(_UpperCAmelCase					,    _UpperCAmelCase     )     )
										else:
													self.assertEqual(_UpperCAmelCase					,    _UpperCAmelCase     )
				@slow
				def 			a_			(self     )      ->			List[str]:
							self.check_tokenizer_from_pretrained(_UpperCAmelCase     )
				def 			a_			(self     )      ->			List[Any]:
							self.assertEqual(TensorType("tf"     )					,    TensorType.TENSORFLOW     )
							self.assertEqual(TensorType("pt"     )					,    TensorType.PYTORCH     )
							self.assertEqual(TensorType("np"     )					,    TensorType.NUMPY     )
				@require_tokenizers
				def 			a_			(self     )      ->			int:
							import numpy as np

							__UpperCamelCase :     Optional[int]										=  BertTokenizer.from_pretrained("bert-base-cased"     )
							__UpperCamelCase :     int										=  BertTokenizerFast.from_pretrained("bert-base-cased"     )

							# Python no tensor
							with self.subTest("BatchEncoding (Python, return_tensors=None)"     ):
										self.assert_dump_and_restore(tokenizer_p("Small example to encode"     )     )

							with self.subTest("BatchEncoding (Python, return_tensors=NUMPY)"     ):
										self.assert_dump_and_restore(
										    tokenizer_p("Small example to encode"					,    return_tensors=TensorType.NUMPY     )					,    np.array_equal     )

							with self.subTest("BatchEncoding (Rust, return_tensors=None)"     ):
										self.assert_dump_and_restore(tokenizer_r("Small example to encode"     )     )

							with self.subTest("BatchEncoding (Rust, return_tensors=NUMPY)"     ):
										self.assert_dump_and_restore(
										    tokenizer_r("Small example to encode"					,    return_tensors=TensorType.NUMPY     )					,    np.array_equal     )
				@require_tf
				@require_tokenizers
				def 			a_			(self     )      ->			str:
							import tensorflow as tf

							def tf_array_equals(_UpperCAmelCase					,    _UpperCAmelCase     ):
										return tf.reduce_all(tf.equal(_UpperCAmelCase					,    _UpperCAmelCase     )     )

							__UpperCamelCase :     Optional[Any]										=  BertTokenizer.from_pretrained("bert-base-cased"     )
							__UpperCamelCase :     int										=  BertTokenizerFast.from_pretrained("bert-base-cased"     )

							with self.subTest("BatchEncoding (Python, return_tensors=TENSORFLOW)"     ):
										self.assert_dump_and_restore(
										    tokenizer_p("Small example to encode"					,    return_tensors=TensorType.TENSORFLOW     )					,    _UpperCAmelCase     )

							with self.subTest("BatchEncoding (Rust, return_tensors=TENSORFLOW)"     ):
										self.assert_dump_and_restore(
										    tokenizer_r("Small example to encode"					,    return_tensors=TensorType.TENSORFLOW     )					,    _UpperCAmelCase     )
				@require_torch
				@require_tokenizers
				def 			a_			(self     )      ->			int:
							import torch

							__UpperCamelCase :     Any										=  BertTokenizer.from_pretrained("bert-base-cased"     )
							__UpperCamelCase :     Any										=  BertTokenizerFast.from_pretrained("bert-base-cased"     )

							with self.subTest("BatchEncoding (Python, return_tensors=PYTORCH)"     ):
										self.assert_dump_and_restore(
										    tokenizer_p("Small example to encode"					,    return_tensors=TensorType.PYTORCH     )					,    torch.equal     )

							with self.subTest("BatchEncoding (Rust, return_tensors=PYTORCH)"     ):
										self.assert_dump_and_restore(
										    tokenizer_r("Small example to encode"					,    return_tensors=TensorType.PYTORCH     )					,    torch.equal     )
				@require_tokenizers
				def 			a_			(self     )      ->			List[str]:
							__UpperCamelCase :     Union[str, Any]										=  BertTokenizer.from_pretrained("bert-base-cased"     )
							__UpperCamelCase :     List[str]										=  BertTokenizerFast.from_pretrained("bert-base-cased"     )

							with self.subTest("Python Tokenizer"     ):
										self.assertFalse(tokenizer_p("Small example to_encode"     ).is_fast     )

							with self.subTest("Rust Tokenizer"     ):
										self.assertTrue(tokenizer_r("Small example to_encode"     ).is_fast     )
				@require_tokenizers
				def 			a_			(self     )      ->			Optional[int]:
							__UpperCamelCase :     Tuple										=  BertTokenizerFast.from_pretrained("bert-base-cased"     )
							__UpperCamelCase :     Dict										=  tokenizer_r(["Test", "\xad", "test"]					,    is_split_into_words=_UpperCAmelCase     )

							self.assertEqual(encoded.word_to_tokens(0     )					,    TokenSpan(start=1					,    end=2     )     )
							self.assertEqual(encoded.word_to_tokens(1     )					,    _UpperCAmelCase     )
							self.assertEqual(encoded.word_to_tokens(2     )					,    TokenSpan(start=2					,    end=3     )     )
				def 			a_			(self     )      ->			int:
							__UpperCamelCase :     Optional[Any]										=  BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]}     )
							__UpperCamelCase :     List[str]										=  batch.convert_to_tensors(tensor_type="np"     )
							self.assertEqual(tensor_batch["inputs"].shape					,    (2, 3)     )
							self.assertEqual(tensor_batch["labels"].shape					,    (2,)     )
							# test converting the converted
							with CaptureStderr() as cs:
										__UpperCamelCase :     Optional[int]										=  batch.convert_to_tensors(tensor_type="np"     )
							self.assertFalse(len(cs.err     )					,    msg=f"should have no warning, but got {cs.err}"     )

							__UpperCamelCase :     Union[str, Any]										=  BatchEncoding({"inputs": [1, 2, 3], "labels": 0}     )
							__UpperCamelCase :     Union[str, Any]										=  batch.convert_to_tensors(tensor_type="np"					,    prepend_batch_axis=_UpperCAmelCase     )
							self.assertEqual(tensor_batch["inputs"].shape					,    (1, 3)     )
							self.assertEqual(tensor_batch["labels"].shape					,    (1,)     )
				@require_torch
				def 			a_			(self     )      ->			Union[str, Any]:
							__UpperCamelCase :     List[str]										=  BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]}     )
							__UpperCamelCase :     Dict										=  batch.convert_to_tensors(tensor_type="pt"     )
							self.assertEqual(tensor_batch["inputs"].shape					,    (2, 3)     )
							self.assertEqual(tensor_batch["labels"].shape					,    (2,)     )
							# test converting the converted
							with CaptureStderr() as cs:
										__UpperCamelCase :     Tuple										=  batch.convert_to_tensors(tensor_type="pt"     )
							self.assertFalse(len(cs.err     )					,    msg=f"should have no warning, but got {cs.err}"     )

							__UpperCamelCase :     int										=  BatchEncoding({"inputs": [1, 2, 3], "labels": 0}     )
							__UpperCamelCase :     Any										=  batch.convert_to_tensors(tensor_type="pt"					,    prepend_batch_axis=_UpperCAmelCase     )
							self.assertEqual(tensor_batch["inputs"].shape					,    (1, 3)     )
							self.assertEqual(tensor_batch["labels"].shape					,    (1,)     )
				@require_tf
				def 			a_			(self     )      ->			int:
							__UpperCamelCase :     List[str]										=  BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]}     )
							__UpperCamelCase :     int										=  batch.convert_to_tensors(tensor_type="tf"     )
							self.assertEqual(tensor_batch["inputs"].shape					,    (2, 3)     )
							self.assertEqual(tensor_batch["labels"].shape					,    (2,)     )
							# test converting the converted
							with CaptureStderr() as cs:
										__UpperCamelCase :     Tuple										=  batch.convert_to_tensors(tensor_type="tf"     )
							self.assertFalse(len(cs.err     )					,    msg=f"should have no warning, but got {cs.err}"     )

							__UpperCamelCase :     Optional[int]										=  BatchEncoding({"inputs": [1, 2, 3], "labels": 0}     )
							__UpperCamelCase :     Union[str, Any]										=  batch.convert_to_tensors(tensor_type="tf"					,    prepend_batch_axis=_UpperCAmelCase     )
							self.assertEqual(tensor_batch["inputs"].shape					,    (1, 3)     )
							self.assertEqual(tensor_batch["labels"].shape					,    (1,)     )
				@require_flax
				def 			a_			(self     )      ->			List[Any]:
							__UpperCamelCase :     int										=  BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]}     )
							__UpperCamelCase :     Dict										=  batch.convert_to_tensors(tensor_type="jax"     )
							self.assertEqual(tensor_batch["inputs"].shape					,    (2, 3)     )
							self.assertEqual(tensor_batch["labels"].shape					,    (2,)     )
							# test converting the converted
							with CaptureStderr() as cs:
										__UpperCamelCase :     Union[str, Any]										=  batch.convert_to_tensors(tensor_type="jax"     )
							self.assertFalse(len(cs.err     )					,    msg=f"should have no warning, but got {cs.err}"     )

							__UpperCamelCase :     Tuple										=  BatchEncoding({"inputs": [1, 2, 3], "labels": 0}     )
							__UpperCamelCase :     List[str]										=  batch.convert_to_tensors(tensor_type="jax"					,    prepend_batch_axis=_UpperCAmelCase     )
							self.assertEqual(tensor_batch["inputs"].shape					,    (1, 3)     )
							self.assertEqual(tensor_batch["labels"].shape					,    (1,)     )
				def 			a_			(self     )      ->			Optional[int]:
							__UpperCamelCase :     int										=  [{"input_ids": np.array([0, 1, 2]     )}, {"input_ids": np.array([0, 1, 2, 3]     )}]
							__UpperCamelCase :     Optional[int]										=  BertTokenizer.from_pretrained("bert-base-cased"     )

							__UpperCamelCase :     Any										=  tokenizer.pad(_UpperCAmelCase					,    padding=_UpperCAmelCase     )
							self.assertTrue(isinstance(batch["input_ids"]					,    np.ndarray     )     )
							self.assertEqual(batch["input_ids"].tolist()					,    [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )
							__UpperCamelCase :     str										=  tokenizer.pad(_UpperCAmelCase					,    padding=_UpperCAmelCase					,    return_tensors="np"     )
							self.assertTrue(isinstance(batch["input_ids"]					,    np.ndarray     )     )
							self.assertEqual(batch["input_ids"].tolist()					,    [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )
				@require_torch
				def 			a_			(self     )      ->			str:
							import torch

							__UpperCamelCase :     Any										=  [{"input_ids": torch.tensor([0, 1, 2]     )}, {"input_ids": torch.tensor([0, 1, 2, 3]     )}]
							__UpperCamelCase :     str										=  BertTokenizer.from_pretrained("bert-base-cased"     )

							__UpperCamelCase :     Dict										=  tokenizer.pad(_UpperCAmelCase					,    padding=_UpperCAmelCase     )
							self.assertTrue(isinstance(batch["input_ids"]					,    torch.Tensor     )     )
							self.assertEqual(batch["input_ids"].tolist()					,    [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )
							__UpperCamelCase :     Dict										=  tokenizer.pad(_UpperCAmelCase					,    padding=_UpperCAmelCase					,    return_tensors="pt"     )
							self.assertTrue(isinstance(batch["input_ids"]					,    torch.Tensor     )     )
							self.assertEqual(batch["input_ids"].tolist()					,    [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )
				@require_tf
				def 			a_			(self     )      ->			Optional[Any]:
							import tensorflow as tf

							__UpperCamelCase :     Any										=  [{"input_ids": tf.constant([0, 1, 2]     )}, {"input_ids": tf.constant([0, 1, 2, 3]     )}]
							__UpperCamelCase :     List[Any]										=  BertTokenizer.from_pretrained("bert-base-cased"     )

							__UpperCamelCase :     int										=  tokenizer.pad(_UpperCAmelCase					,    padding=_UpperCAmelCase     )
							self.assertTrue(isinstance(batch["input_ids"]					,    tf.Tensor     )     )
							self.assertEqual(batch["input_ids"].numpy().tolist()					,    [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )
							__UpperCamelCase :     int										=  tokenizer.pad(_UpperCAmelCase					,    padding=_UpperCAmelCase					,    return_tensors="tf"     )
							self.assertTrue(isinstance(batch["input_ids"]					,    tf.Tensor     )     )
							self.assertEqual(batch["input_ids"].numpy().tolist()					,    [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]     )
				@require_tokenizers
				def 			a_			(self     )      ->			Dict:
							__UpperCamelCase :     Optional[Any]										=  Tokenizer(WordPiece(unk_token="[UNK]"     )     )
							PreTrainedTokenizerFast(tokenizer_object=_UpperCAmelCase     )





				@require_tokenizers
				def 			a_			(self     )      ->			Any:
							__UpperCamelCase :     Optional[Any]										=  Tokenizer(WordPiece(unk_token="[UNK]"     )     )
							with tempfile.TemporaryDirectory() as tmpdirname:
										bert_tokenizer.save(os.path.join(_UpperCAmelCase					,    "tokenizer.json"     )     )
										PreTrainedTokenizerFast(tokenizer_file=os.path.join(_UpperCAmelCase					,    "tokenizer.json"     )     )









'''simple docstring'''






import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import (
    TaBlock,
    TaConfig,
    TaLayerNorm,
)

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin



class    A						(					SCREAMING_SNAKE_CASE__     ,				SCREAMING_SNAKE_CASE__     ,				SCREAMING_SNAKE_CASE__						):




				'''simple docstring'''





				@register_to_config
				def __init__(self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase = False					,    )      ->			Tuple:
							super().__init__()

							__UpperCamelCase :     Optional[int]										=  nn.Linear(_UpperCAmelCase					,    _UpperCAmelCase					,    bias=_UpperCAmelCase     )

							__UpperCamelCase :     str										=  nn.Embedding(_UpperCAmelCase					,    _UpperCAmelCase     )
							__UpperCamelCase :     Any										=  False

							__UpperCamelCase :     Union[str, Any]										=  nn.Dropout(p=_UpperCAmelCase     )

							__UpperCamelCase :     Dict										=  TaConfig(
							    d_model=_UpperCAmelCase					,    num_heads=_UpperCAmelCase					,    d_kv=_UpperCAmelCase					,    d_ff=_UpperCAmelCase					,    feed_forward_proj=_UpperCAmelCase					,    dropout_rate=_UpperCAmelCase					,    is_decoder=_UpperCAmelCase					,    is_encoder_decoder=_UpperCAmelCase					,    )
							__UpperCamelCase :     Tuple										=  nn.ModuleList()
							for lyr_num in range(_UpperCAmelCase     ):
										__UpperCamelCase :     List[str]										=  TaBlock(_UpperCAmelCase     )
										self.encoders.append(_UpperCAmelCase     )

							__UpperCamelCase :     Tuple										=  TaLayerNorm(_UpperCAmelCase     )
							__UpperCamelCase :     Any										=  nn.Dropout(p=_UpperCAmelCase     )





				def 			a_			(self					,    _UpperCAmelCase					,    _UpperCAmelCase     )      ->			int:
							__UpperCamelCase :     Any										=  self.input_proj(_UpperCAmelCase     )

							# terminal relative positional encodings
							__UpperCamelCase :     Optional[Any]										=  encoder_inputs.shape[1]
							__UpperCamelCase :     Optional[int]										=  torch.arange(_UpperCAmelCase					,    device=encoder_inputs.device     )

							__UpperCamelCase :     Tuple										=  encoder_inputs_mask.sum(-1     )
							__UpperCamelCase :     Optional[Any]										=  torch.roll(input_positions.unsqueeze(0     )					,    tuple(seq_lens.tolist()     )					,    dims=0     )
							x += self.position_encoding(_UpperCAmelCase     )

							__UpperCamelCase :     Any										=  self.dropout_pre(_UpperCAmelCase     )

							# inverted the attention mask
							__UpperCamelCase :     Union[str, Any]										=  encoder_inputs.size()
							__UpperCamelCase :     Optional[Any]										=  self.get_extended_attention_mask(_UpperCAmelCase					,    _UpperCAmelCase     )

							for lyr in self.encoders:
										__UpperCamelCase :     List[Any]										=  lyr(_UpperCAmelCase					,    _UpperCAmelCase     )[0]
							__UpperCamelCase :     Union[str, Any]										=  self.layer_norm(_UpperCAmelCase     )

							return self.dropout_post(_UpperCAmelCase     ), encoder_inputs_mask



