



'''simple docstring'''





import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import (
     MODEL_MAPPING,
     BeitForImageClassification,
     BeitForMaskedImageModeling,
     BeitForSemanticSegmentation,
     BeitModel,
 )
 from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 import PIL
 from PIL import Image

 from transformers import BeitImageProcessor


class 				snake_case :



       def __init__(				self     ,UpperCAmelCase_     ,UpperCAmelCase_=100     ,UpperCAmelCase_=13     ,UpperCAmelCase_=30     ,UpperCAmelCase_=2     ,UpperCAmelCase_=3     ,UpperCAmelCase_=True     ,UpperCAmelCase_=True     ,UpperCAmelCase_=32     ,UpperCAmelCase_=4     ,UpperCAmelCase_=4     ,UpperCAmelCase_=37     ,UpperCAmelCase_="gelu"     ,UpperCAmelCase_=0.1     ,UpperCAmelCase_=0.1     ,UpperCAmelCase_=10     ,UpperCAmelCase_=0.02     ,UpperCAmelCase_=3     ,UpperCAmelCase_=None     ,UpperCAmelCase_=[0, 1, 2, 3]     ,)  ->							Any:
           lowercase__		       =						parent
           lowercase__		       =						100
           lowercase__		       =						batch_size
           lowercase__		       =						image_size
           lowercase__		       =						patch_size
           lowercase__		       =						num_channels
           lowercase__		       =						is_training
           lowercase__		       =						use_labels
           lowercase__		       =						hidden_size
           lowercase__		       =						num_hidden_layers
           lowercase__		       =						num_attention_heads
           lowercase__		       =						intermediate_size
           lowercase__		       =						hidden_act
           lowercase__		       =						hidden_dropout_prob
           lowercase__		       =						attention_probs_dropout_prob
           lowercase__		       =						type_sequence_label_size
           lowercase__		       =						initializer_range
           lowercase__		       =						scope
           lowercase__		       =						out_indices
           lowercase__		       =						num_labels

           # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
           lowercase__		       =						(image_size // patch_size) ** 2
           lowercase__		       =						num_patches + 1



       def 							_a      (				self      )  ->							str:
           lowercase__		       =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

           lowercase__		       =						None
           lowercase__		       =						None
           if self.use_labels:
               lowercase__		       =						ids_tensor([self.batch_size]     ,self.type_sequence_label_size      )
               lowercase__		       =						ids_tensor([self.batch_size, self.image_size, self.image_size]     ,self.num_labels      )

           lowercase__		       =						self.get_config()

           return config, pixel_values, labels, pixel_labels



       def 							_a      (				self      )  ->							List[str]:
           return BeitConfig(
               vocab_size=self.vocab_size     ,image_size=self.image_size     ,patch_size=self.patch_size     ,num_channels=self.num_channels     ,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     ,is_decoder=UpperCAmelCase_     ,initializer_range=self.initializer_range     ,out_indices=self.out_indices     ,)



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Optional[int]:
           lowercase__		       =						BeitModel(config=UpperCAmelCase_      )
           model.to(UpperCAmelCase_      )
           model.eval()
           lowercase__		       =						model(UpperCAmelCase_      )
           self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, self.seq_length, self.hidden_size)      )



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Tuple:
           lowercase__		       =						BeitForMaskedImageModeling(config=UpperCAmelCase_      )
           model.to(UpperCAmelCase_      )
           model.eval()
           lowercase__		       =						model(UpperCAmelCase_      )
           self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.seq_length - 1, self.vocab_size)      )



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Dict:
           lowercase__		       =						self.type_sequence_label_size
           lowercase__		       =						BeitForImageClassification(UpperCAmelCase_      )
           model.to(UpperCAmelCase_      )
           model.eval()
           lowercase__		       =						model(UpperCAmelCase_     ,labels=UpperCAmelCase_      )
           self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.type_sequence_label_size)      )

           # test greyscale images
           lowercase__		       =						1
           lowercase__		       =						BeitForImageClassification(UpperCAmelCase_      )
           model.to(UpperCAmelCase_      )
           model.eval()

           lowercase__		       =						floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
           lowercase__		       =						model(UpperCAmelCase_     ,labels=UpperCAmelCase_      )
           self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.type_sequence_label_size)      )



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Dict:
           lowercase__		       =						self.num_labels
           lowercase__		       =						BeitForSemanticSegmentation(UpperCAmelCase_      )
           model.to(UpperCAmelCase_      )
           model.eval()
           lowercase__		       =						model(UpperCAmelCase_      )
           self.parent.assertEqual(
               result.logits.shape     ,(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)      )
           lowercase__		       =						model(UpperCAmelCase_     ,labels=UpperCAmelCase_      )
           self.parent.assertEqual(
               result.logits.shape     ,(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)      )







       def 							_a      (				self      )  ->							str:
           lowercase__		       =						self.prepare_config_and_inputs()
           lowercase__       ,       lowercase__       ,       lowercase__       ,       lowercase__		       =						config_and_inputs
           lowercase__		       =						{"pixel_values": pixel_values}
           return config, inputs_dict





@require_torch
class 				snake_case (UpperCamelCase		,      UpperCamelCase		,      unittest.TestCase   ):
       lowerCAmelCase__  :Optional[Any]						   =					(
           (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
           if is_torch_available()
           else ()
       )
       lowerCAmelCase__  :Optional[int]						   =					(
           {
               "feature-extraction": BeitModel,
               "image-classification": BeitForImageClassification,
               "image-segmentation": BeitForSemanticSegmentation,
           }
           if is_torch_available()
           else {}
       )

       lowerCAmelCase__  :Any						   =					False
       lowerCAmelCase__  :Optional[Any]						   =					False
       lowerCAmelCase__  :Optional[Any]						   =					False



       def 							_a      (				self      )  ->							Union[str, Any]:
           lowercase__		       =						BeitModelTester(self      )
           lowercase__		       =						ConfigTester(self     ,config_class=UpperCAmelCase_     ,has_text_modality=UpperCAmelCase_     ,hidden_size=37      )



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



       @unittest.skip(reason="BEiT does not use inputs_embeds"      )
       def 							_a      (				self      )  ->							Dict:
           pass



       @require_torch_multi_gpu
       @unittest.skip(reason="BEiT has some layers using `add_module` which doesn't work well with `nn.DataParallel`"      )
       def 							_a      (				self      )  ->							int:
           pass



       def 							_a      (				self      )  ->							Any:
           lowercase__       ,       lowercase__		       =						self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               lowercase__		       =						model_class(UpperCAmelCase_      )
               self.assertIsInstance(model.get_input_embeddings()     ,(nn.Module)      )
               lowercase__		       =						model.get_output_embeddings()
               self.assertTrue(x is None or isinstance(UpperCAmelCase_     ,nn.Linear      )      )



       def 							_a      (				self      )  ->							Tuple:
           lowercase__       ,       lowercase__		       =						self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               lowercase__		       =						model_class(UpperCAmelCase_      )
               lowercase__		       =						inspect.signature(model.forward      )
               # signature.parameters is an OrderedDict => so arg_names order is deterministic
               lowercase__		       =						[*signature.parameters.keys()]

               lowercase__		       =						["pixel_values"]
               self.assertListEqual(arg_names[:1]     ,UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Tuple:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							List[Any]:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Any:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							str:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_semantic_segmentation(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							int:
           if not self.model_tester.is_training:
               return

           lowercase__       ,       lowercase__		       =						self.model_tester.prepare_config_and_inputs_for_common()
           lowercase__		       =						True

           for model_class in self.all_model_classes:
               # we don't test BeitForMaskedImageModeling
               if model_class in [*get_values(UpperCAmelCase_      ), BeitForMaskedImageModeling]:
                   continue

               lowercase__		       =						model_class(UpperCAmelCase_      )
               model.to(UpperCAmelCase_      )
               model.train()
               lowercase__		       =						self._prepare_for_class(UpperCAmelCase_     ,UpperCAmelCase_     ,return_labels=UpperCAmelCase_      )
               lowercase__		       =						model(**UpperCAmelCase_      ).loss
               loss.backward()



       def 							_a      (				self      )  ->							str:
           lowercase__       ,       lowercase__		       =						self.model_tester.prepare_config_and_inputs_for_common()
           if not self.model_tester.is_training:
               return

           lowercase__		       =						False
           lowercase__		       =						True

           for model_class in self.all_model_classes:
               # we don't test BeitForMaskedImageModeling
               if (
                   model_class in [*get_values(UpperCAmelCase_      ), BeitForMaskedImageModeling]
                   or not model_class.supports_gradient_checkpointing
               ):
                   continue

               lowercase__		       =						model_class(UpperCAmelCase_      )
               model.gradient_checkpointing_enable()
               model.to(UpperCAmelCase_      )
               model.train()
               lowercase__		       =						self._prepare_for_class(UpperCAmelCase_     ,UpperCAmelCase_     ,return_labels=UpperCAmelCase_      )
               lowercase__		       =						model(**UpperCAmelCase_      ).loss
               loss.backward()



       def 							_a      (				self      )  ->							int:
           lowercase__       ,       lowercase__		       =						self.model_tester.prepare_config_and_inputs_for_common()

           lowercase__		       =						_config_zero_init(UpperCAmelCase_      )
           for model_class in self.all_model_classes:
               lowercase__		       =						model_class(config=UpperCAmelCase_      )
               for name, param in model.named_parameters():
                   # we skip lambda parameters as these require special initial values
                   # determined by config.layer_scale_init_value
                   if "lambda" in name:
                       continue
                   if param.requires_grad:
                       self.assertIn(
                           ((param.data.mean() * 1E9).round() / 1E9).item()     ,[0.0, 1.0]     ,msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''     ,)







       @slow
       def 							_a      (				self      )  ->							Optional[Any]:
           for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               lowercase__		       =						BeitModel.from_pretrained(UpperCAmelCase_      )
               self.assertIsNotNone(UpperCAmelCase_      )





def        lowerCamelCase  (				):





    '''simple docstring'''


    lowercase__		       =						Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"       )
    return image



@require_torch
@require_vision
class 				snake_case (unittest.TestCase   ):



       @cached_property
       def 							_a      (				self      )  ->							Any:
           return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"      ) if is_vision_available() else None



       @slow
       def 							_a      (				self      )  ->							int:
           lowercase__		       =						BeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"      ).to(UpperCAmelCase_      )

           lowercase__		       =						self.default_image_processor
           lowercase__		       =						prepare_img()
           lowercase__		       =						image_processor(images=UpperCAmelCase_     ,return_tensors="pt"      ).pixel_values.to(UpperCAmelCase_      )

           # prepare bool_masked_pos
           lowercase__		       =						torch.ones((1, 196)     ,dtype=torch.bool      ).to(UpperCAmelCase_      )

           # forward pass
           with torch.no_grad():
               lowercase__		       =						model(pixel_values=UpperCAmelCase_     ,bool_masked_pos=UpperCAmelCase_      )
           lowercase__		       =						outputs.logits

           # verify the logits
           lowercase__		       =						torch.Size((1, 196, 8_192)      )
           self.assertEqual(logits.shape     ,UpperCAmelCase_      )

           lowercase__		       =						torch.tensor(
               [[-3.24_37, 0.50_72, -13.91_74], [-3.24_56, 0.49_48, -13.94_01], [-3.20_33, 0.51_21, -13.85_50]]      ).to(UpperCAmelCase_      )

           self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]     ,UpperCAmelCase_     ,atol=1E-2      )      )



       @slow
       def 							_a      (				self      )  ->							Optional[Any]:
           lowercase__		       =						BeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"      ).to(UpperCAmelCase_      )

           lowercase__		       =						self.default_image_processor
           lowercase__		       =						prepare_img()
           lowercase__		       =						image_processor(images=UpperCAmelCase_     ,return_tensors="pt"      ).to(UpperCAmelCase_      )

           # forward pass
           with torch.no_grad():
               lowercase__		       =						model(**UpperCAmelCase_      )
           lowercase__		       =						outputs.logits

           # verify the logits
           lowercase__		       =						torch.Size((1, 1_000)      )
           self.assertEqual(logits.shape     ,UpperCAmelCase_      )

           lowercase__		       =						torch.tensor([-1.23_85, -1.09_87, -1.01_08]      ).to(UpperCAmelCase_      )

           self.assertTrue(torch.allclose(logits[0, :3]     ,UpperCAmelCase_     ,atol=1E-4      )      )

           lowercase__		       =						281
           self.assertEqual(logits.argmax(-1      ).item()     ,UpperCAmelCase_      )



       @slow
       def 							_a      (				self      )  ->							Union[str, Any]:
           lowercase__		       =						BeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"      ).to(
               UpperCAmelCase_      )

           lowercase__		       =						self.default_image_processor
           lowercase__		       =						prepare_img()
           lowercase__		       =						image_processor(images=UpperCAmelCase_     ,return_tensors="pt"      ).to(UpperCAmelCase_      )

           # forward pass
           with torch.no_grad():
               lowercase__		       =						model(**UpperCAmelCase_      )
           lowercase__		       =						outputs.logits

           # verify the logits
           lowercase__		       =						torch.Size((1, 21_841)      )
           self.assertEqual(logits.shape     ,UpperCAmelCase_      )

           lowercase__		       =						torch.tensor([1.68_81, -0.27_87, 0.59_01]      ).to(UpperCAmelCase_      )

           self.assertTrue(torch.allclose(logits[0, :3]     ,UpperCAmelCase_     ,atol=1E-4      )      )

           lowercase__		       =						2_396
           self.assertEqual(logits.argmax(-1      ).item()     ,UpperCAmelCase_      )



       @slow
       def 							_a      (				self      )  ->							Union[str, Any]:
           lowercase__		       =						BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"      )
           lowercase__		       =						model.to(UpperCAmelCase_      )

           lowercase__		       =						BeitImageProcessor(do_resize=UpperCAmelCase_     ,size=640     ,do_center_crop=UpperCAmelCase_      )

           lowercase__		       =						load_dataset("hf-internal-testing/fixtures_ade20k"     ,split="test"      )
           lowercase__		       =						Image.open(ds[0]["file"]      )
           lowercase__		       =						image_processor(images=UpperCAmelCase_     ,return_tensors="pt"      ).to(UpperCAmelCase_      )

           # forward pass
           with torch.no_grad():
               lowercase__		       =						model(**UpperCAmelCase_      )
           lowercase__		       =						outputs.logits

           # verify the logits
           lowercase__		       =						torch.Size((1, 150, 160, 160)      )
           self.assertEqual(logits.shape     ,UpperCAmelCase_      )

           lowercase__		       =						version.parse(PIL.__version__      ) < version.parse("9.0.0"      )

           if is_pillow_less_than_a:
               lowercase__		       =						torch.tensor(
                   [
                       [[-4.92_25, -2.39_54, -3.05_22], [-2.88_22, -1.00_46, -1.75_61], [-2.95_49, -1.32_28, -2.13_47]],
                       [[-5.81_68, -3.41_29, -4.07_78], [-3.86_51, -2.22_14, -3.02_77], [-3.83_56, -2.46_43, -3.35_35]],
                       [[-0.00_78, 3.99_52, 4.07_54], [2.98_56, 4.69_44, 5.00_35], [3.24_13, 4.78_13, 4.99_69]],
                   ]     ,device=UpperCAmelCase_     ,)
           else:
               lowercase__		       =						torch.tensor(
                   [
                       [[-4.89_60, -2.36_88, -3.03_55], [-2.84_78, -0.98_36, -1.74_18], [-2.94_49, -1.33_32, -2.14_56]],
                       [[-5.80_81, -3.41_24, -4.10_06], [-3.85_61, -2.20_81, -3.03_23], [-3.83_65, -2.46_01, -3.36_69]],
                       [[-0.03_09, 3.98_68, 4.05_40], [2.96_40, 4.68_77, 4.99_76], [3.20_81, 4.76_90, 4.99_42]],
                   ]     ,device=UpperCAmelCase_     ,)

           self.assertTrue(torch.allclose(logits[0, :3, :3, :3]     ,UpperCAmelCase_     ,atol=1E-4      )      )







       @slow
       def 							_a      (				self      )  ->							str:
           lowercase__		       =						BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"      )
           lowercase__		       =						model.to(UpperCAmelCase_      )

           lowercase__		       =						BeitImageProcessor(do_resize=UpperCAmelCase_     ,size=640     ,do_center_crop=UpperCAmelCase_      )

           lowercase__		       =						load_dataset("hf-internal-testing/fixtures_ade20k"     ,split="test"      )
           lowercase__		       =						Image.open(ds[0]["file"]      )
           lowercase__		       =						image_processor(images=UpperCAmelCase_     ,return_tensors="pt"      ).to(UpperCAmelCase_      )

           # forward pass
           with torch.no_grad():
               lowercase__		       =						model(**UpperCAmelCase_      )

           lowercase__		       =						outputs.logits.detach().cpu()

           lowercase__		       =						image_processor.post_process_semantic_segmentation(outputs=UpperCAmelCase_     ,target_sizes=[(500, 300)]      )
           lowercase__		       =						torch.Size((500, 300)      )
           self.assertEqual(segmentation[0].shape     ,UpperCAmelCase_      )

           lowercase__		       =						image_processor.post_process_semantic_segmentation(outputs=UpperCAmelCase_      )
           lowercase__		       =						torch.Size((160, 160)      )
           self.assertEqual(segmentation[0].shape     ,UpperCAmelCase_      )









'''simple docstring'''





import numpy as np
import torch
from torch.utils.data import Dataset

from utils import logger


class 				snake_case (UpperCamelCase   ):



       def __init__(				self     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							int:
           lowercase__		       =						params

           lowercase__		       =						np.array(UpperCAmelCase_      )
           lowercase__		       =						np.array([len(UpperCAmelCase_      ) for t in data]      )

           self.check()
           self.remove_long_sequences()
           self.remove_empty_sequences()
           self.remove_unknown_sequences()
           self.check()
           self.print_statistics()



       def __getitem__(				self     ,UpperCAmelCase_      )  ->							Optional[int]:
           return (self.token_ids[index], self.lengths[index])



       def __len__(				self      )  ->							Union[str, Any]:
           return len(self.lengths      )



       def 							_a      (				self      )  ->							int:
           assert len(self.token_ids      ) == len(self.lengths      )
           assert all(self.lengths[i] == len(self.token_ids[i]      ) for i in range(len(self.lengths      )      )      )



       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						self.params.max_model_input_size
           lowercase__		       =						self.lengths > max_len
           logger.info(F'''Splitting {sum(UpperCAmelCase_      )} too long sequences.'''      )

           def divide_chunks(UpperCAmelCase_     ,UpperCAmelCase_      ):
               return [l[i : i + n] for i in range(0     ,len(UpperCAmelCase_      )     ,UpperCAmelCase_      )]

           lowercase__		       =						[]
           lowercase__		       =						[]
           if self.params.mlm:
               lowercase__       ,       lowercase__		       =						self.params.special_tok_ids["cls_token"], self.params.special_tok_ids["sep_token"]
           else:
               lowercase__       ,       lowercase__		       =						self.params.special_tok_ids["bos_token"], self.params.special_tok_ids["eos_token"]

           for seq_, len_ in zip(self.token_ids     ,self.lengths      ):
               assert (seq_[0] == cls_id) and (seq_[-1] == sep_id), seq_
               if len_ <= max_len:
                   new_tok_ids.append(seq_      )
                   new_lengths.append(len_      )
               else:
                   lowercase__		       =						[]
                   for sub_s in divide_chunks(seq_     ,max_len - 2      ):
                       if sub_s[0] != cls_id:
                           lowercase__		       =						np.insert(UpperCAmelCase_     ,0     ,UpperCAmelCase_      )
                       if sub_s[-1] != sep_id:
                           lowercase__		       =						np.insert(UpperCAmelCase_     ,len(UpperCAmelCase_      )     ,UpperCAmelCase_      )
                       assert len(UpperCAmelCase_      ) <= max_len
                       assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
                       sub_seqs.append(UpperCAmelCase_      )

                   new_tok_ids.extend(UpperCAmelCase_      )
                   new_lengths.extend([len(UpperCAmelCase_      ) for l in sub_seqs]      )

           lowercase__		       =						np.array(UpperCAmelCase_      )
           lowercase__		       =						np.array(UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Any:
           lowercase__		       =						len(self      )
           lowercase__		       =						self.lengths > 11
           lowercase__		       =						self.token_ids[indices]
           lowercase__		       =						self.lengths[indices]
           lowercase__		       =						len(self      )
           logger.info(F'''Remove {init_size - new_size} too short (<=11 tokens) sequences.'''      )



       def 							_a      (				self      )  ->							List[Any]:
           if "unk_token" not in self.params.special_tok_ids:
               return
           else:
               lowercase__		       =						self.params.special_tok_ids["unk_token"]
           lowercase__		       =						len(self      )
           lowercase__		       =						np.array([np.count_nonzero(a == unk_token_id      ) for a in self.token_ids]      )
           lowercase__		       =						(unk_occs / self.lengths) < 0.5
           lowercase__		       =						self.token_ids[indices]
           lowercase__		       =						self.lengths[indices]
           lowercase__		       =						len(self      )
           logger.info(F'''Remove {init_size - new_size} sequences with a high level of unknown tokens (50%).'''      )



       def 							_a      (				self      )  ->							Optional[int]:
           if not self.params.is_master:
               return
           logger.info(F'''{len(self      )} sequences'''      )
           # data_len = sum(self.lengths)
           # nb_unique_tokens = len(Counter(list(chain(*self.token_ids))))
           # logger.info(f'{data_len} tokens ({nb_unique_tokens} unique)')

           # unk_idx = self.params.special_tok_ids['unk_token']
           # nb_unknown = sum([(t==unk_idx).sum() for t in self.token_ids])
           # logger.info(f'{nb_unknown} unknown tokens (covering {100*nb_unknown/data_len:.2f}% of the data)')







       def 							_a      (				self     ,UpperCAmelCase_      )  ->							List[str]:
           lowercase__		       =						[t[0] for t in batch]
           lowercase__		       =						[t[1] for t in batch]
           assert len(UpperCAmelCase_      ) == len(UpperCAmelCase_      )

           # Max for paddings
           lowercase__		       =						max(UpperCAmelCase_      )

           # Pad token ids
           if self.params.mlm:
               lowercase__		       =						self.params.special_tok_ids["pad_token"]
           else:
               lowercase__		       =						self.params.special_tok_ids["unk_token"]
           lowercase__		       =						[list(t.astype(UpperCAmelCase_      )      ) + [pad_idx] * (max_seq_len_ - len(UpperCAmelCase_      )) for t in token_ids]
           assert len(tk_      ) == len(UpperCAmelCase_      )
           assert all(len(UpperCAmelCase_      ) == max_seq_len_ for t in tk_      )

           lowercase__		       =						torch.tensor(tk_      )  # (bs, max_seq_len_)
           lowercase__		       =						torch.tensor(UpperCAmelCase_      )  # (bs)
           return tk_t, lg_t




