


"""simple docstring"""




import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         FalconForCausalLM,
         FalconForQuestionAnswering,
         FalconForSequenceClassification,
         FalconForTokenClassification,
         FalconModel,
     )
class 						__A							:






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






     def    __A  (		self  ):
       _lowerCAmelCase						:	Dict        =					ids_tensor([self.batch_size, self.seq_length]    ,					self.vocab_size  )

       _lowerCAmelCase						:	Optional[int]        =					None
       if self.use_input_mask:
         _lowerCAmelCase						:	Optional[int]        =					random_attention_mask([self.batch_size, self.seq_length]  )

       _lowerCAmelCase						:	Tuple        =					None

       _lowerCAmelCase						:	List[str]        =					None
       _lowerCAmelCase						:	Any        =					None
       _lowerCAmelCase						:	Any        =					None
       if self.use_labels:
         _lowerCAmelCase						:	str        =					ids_tensor([self.batch_size]    ,					self.type_sequence_label_size  )
         _lowerCAmelCase						:	Any        =					ids_tensor([self.batch_size, self.seq_length]    ,					self.num_labels  )
         _lowerCAmelCase						:	Any        =					ids_tensor([self.batch_size]    ,					self.num_choices  )

       _lowerCAmelCase						:	List[Any]        =					self.get_config()

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






     def    __A  (		self  ):
       return FalconConfig(
           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    ,					is_decoder=a__    ,					initializer_range=self.initializer_range    ,					pad_token_id=1    ,					new_decoder_architecture=a__    ,					)






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__  ):
       _lowerCAmelCase						:	Union[str, Any]        =					FalconModel(config=a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	int        =					model(a__    ,					attention_mask=a__  )
       _lowerCAmelCase						:	Dict        =					model(a__  )
       self.parent.assertEqual(result.last_hidden_state.shape    ,					(self.batch_size, self.seq_length, self.hidden_size)  )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					):
       _lowerCAmelCase						:	Optional[int]        =					True
       _lowerCAmelCase						:	Tuple        =					FalconModel(a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	List[Any]        =					model(
           a__    ,					attention_mask=a__    ,					encoder_hidden_states=a__    ,					encoder_attention_mask=a__    ,					)
       _lowerCAmelCase						:	int        =					model(
           a__    ,					attention_mask=a__    ,					encoder_hidden_states=a__    ,					)
       _lowerCAmelCase						:	Optional[int]        =					model(a__    ,					attention_mask=a__  )
       self.parent.assertEqual(result.last_hidden_state.shape    ,					(self.batch_size, self.seq_length, self.hidden_size)  )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					):
       _lowerCAmelCase						:	Dict        =					FalconForCausalLM(config=a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	Union[str, Any]        =					model(a__    ,					attention_mask=a__    ,					labels=a__  )
       self.parent.assertEqual(result.logits.shape    ,					(self.batch_size, self.seq_length, self.vocab_size)  )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					):
       _lowerCAmelCase						:	Optional[int]        =					True
       _lowerCAmelCase						:	List[str]        =					True
       _lowerCAmelCase						:	List[Any]        =					FalconForCausalLM(config=a__  )
       model.to(a__  )
       model.eval()

       # first forward pass
       _lowerCAmelCase						:	Dict        =					model(
           a__    ,					attention_mask=a__    ,					encoder_hidden_states=a__    ,					encoder_attention_mask=a__    ,					use_cache=a__    ,					)
       _lowerCAmelCase						:	Any        =					outputs.past_key_values

       # create hypothetical multiple next token and extent to next_input_ids
       _lowerCAmelCase						:	Any        =					ids_tensor((self.batch_size, 3)    ,					config.vocab_size  )
       _lowerCAmelCase						:	Optional[int]        =					ids_tensor((self.batch_size, 3)    ,					vocab_size=2  )

       # append to next input_ids and
       _lowerCAmelCase						:	Any        =					torch.cat([input_ids, next_tokens]    ,					dim=-1  )
       _lowerCAmelCase						:	str        =					torch.cat([input_mask, next_mask]    ,					dim=-1  )

       _lowerCAmelCase						:	List[str]        =					model(
           a__    ,					attention_mask=a__    ,					encoder_hidden_states=a__    ,					encoder_attention_mask=a__    ,					output_hidden_states=a__    ,					)["""hidden_states"""][0]
       _lowerCAmelCase						:	Optional[Any]        =					model(
           a__    ,					attention_mask=a__    ,					encoder_hidden_states=a__    ,					encoder_attention_mask=a__    ,					past_key_values=a__    ,					output_hidden_states=a__    ,					)["""hidden_states"""][0]

       # select random slice
       _lowerCAmelCase						:	Optional[int]        =					ids_tensor((1,)    ,					output_from_past.shape[-1]  ).item()
       _lowerCAmelCase						:	Tuple        =					output_from_no_past[:, -3:, random_slice_idx].detach()
       _lowerCAmelCase						:	Union[str, Any]        =					output_from_past[:, :, random_slice_idx].detach()

       self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]  )

       # test that outputs are equal for slice
       self.parent.assertTrue(torch.allclose(a__    ,					a__    ,					atol=1e-3  )  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Union[str, Any]        =					self.prepare_config_and_inputs()
       (
           _lowerCAmelCase     
       )						:	Tuple        =					config_and_inputs
       _lowerCAmelCase						:	Optional[Any]        =					{"""input_ids""": input_ids, """attention_mask""": input_mask}
       return config, inputs_dict







@require_torch
class 						__A							(			SCREAMING_SNAKE_CASE_   ,     SCREAMING_SNAKE_CASE_   ,     SCREAMING_SNAKE_CASE_   ,     unittest.TestCase		):
     _UpperCamelCase							:      Optional[int]		 =		(
         (
             FalconModel,
             FalconForCausalLM,
             FalconForSequenceClassification,
             FalconForTokenClassification,
             FalconForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     _UpperCamelCase							:      int		 =		(FalconForCausalLM,) if is_torch_available() else ()
     _UpperCamelCase							:      Any		 =		(
         {
             "feature-extraction": FalconModel,
             "text-classification": FalconForSequenceClassification,
             "text-generation": FalconForCausalLM,
             "question-answering": FalconForQuestionAnswering,
             "token-classification": FalconForTokenClassification,
             "zero-shot": FalconForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     _UpperCamelCase							:      Tuple		 =		False
     _UpperCamelCase							:      Optional[Any]		 =		False






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[int]        =					FalconModelTester(self  )
       _lowerCAmelCase						:	Dict        =					ConfigTester(self    ,					config_class=a__    ,					hidden_size=37  )






     def    __A  (		self  ):
       self.config_tester.run_common_tests()






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[Any]        =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*a__  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[Any]        =					self.model_tester.prepare_config_and_inputs()
       for alibi in [True, False]:
         _lowerCAmelCase						:	Tuple        =					alibi
         self.model_tester.create_and_check_model(a__    ,					*a__  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[Any]        =					self.model_tester.prepare_config_and_inputs_for_common()
       _lowerCAmelCase						:	List[Any]        =					3
       _lowerCAmelCase						:	Union[str, Any]        =					input_dict["""input_ids"""]
       _lowerCAmelCase						:	Dict        =					input_ids.ne(1  ).to(a__  )
       _lowerCAmelCase						:	Optional[Any]        =					ids_tensor([self.model_tester.batch_size]    ,					self.model_tester.type_sequence_label_size  )
       _lowerCAmelCase						:	Optional[Any]        =					FalconForSequenceClassification(a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	Union[str, Any]        =					model(a__    ,					attention_mask=a__    ,					labels=a__  )
       self.assertEqual(result.logits.shape    ,					(self.model_tester.batch_size, self.model_tester.num_labels)  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Tuple        =					self.model_tester.prepare_config_and_inputs_for_common()
       _lowerCAmelCase						:	Optional[Any]        =					3
       _lowerCAmelCase						:	Optional[int]        =					"""single_label_classification"""
       _lowerCAmelCase						:	List[Any]        =					input_dict["""input_ids"""]
       _lowerCAmelCase						:	Dict        =					input_ids.ne(1  ).to(a__  )
       _lowerCAmelCase						:	List[str]        =					ids_tensor([self.model_tester.batch_size]    ,					self.model_tester.type_sequence_label_size  )
       _lowerCAmelCase						:	Optional[Any]        =					FalconForSequenceClassification(a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	List[str]        =					model(a__    ,					attention_mask=a__    ,					labels=a__  )
       self.assertEqual(result.logits.shape    ,					(self.model_tester.batch_size, self.model_tester.num_labels)  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Union[str, Any]        =					self.model_tester.prepare_config_and_inputs_for_common()
       _lowerCAmelCase						:	List[str]        =					input_dict["""input_ids"""]
       _lowerCAmelCase						:	List[Any]        =					FalconForCausalLM(a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	List[str]        =					model(a__    ,					use_cache=a__  )
       _lowerCAmelCase						:	Dict        =					input_ids.shape[0]
       _lowerCAmelCase						:	Dict        =					model._convert_to_rw_cache(result.past_key_values  )
       _lowerCAmelCase						:	int        =					model._convert_cache_to_standard_format(a__    ,					a__  )
       for layer in range(len(a__  )  ):
         for tensor_idx in range(2  ):
           self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3  )
           self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4  )
           self.assertTrue(
               torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]  )  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Union[str, Any]        =					self.model_tester.prepare_config_and_inputs_for_common()
       _lowerCAmelCase						:	str        =					3
       _lowerCAmelCase						:	List[Any]        =					"""multi_label_classification"""
       _lowerCAmelCase						:	List[Any]        =					input_dict["""input_ids"""]
       _lowerCAmelCase						:	Tuple        =					input_ids.ne(1  ).to(a__  )
       _lowerCAmelCase						:	Optional[int]        =					ids_tensor(
           [self.model_tester.batch_size, config.num_labels]    ,					self.model_tester.type_sequence_label_size  ).to(torch.float  )
       _lowerCAmelCase						:	List[Any]        =					FalconForSequenceClassification(a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	Union[str, Any]        =					model(a__    ,					attention_mask=a__    ,					labels=a__  )
       self.assertEqual(result.logits.shape    ,					(self.model_tester.batch_size, self.model_tester.num_labels)  )






     def    __A  (		self  ):
       # Falcon can have different numbers of KV-heads than the number of query heads, so we need
       # to override this test to use the right head counts.
       for model_class in self.all_generative_model_classes:
         _lowerCAmelCase						:	Union[str, Any]        =					self.model_tester.prepare_config_and_inputs_for_common()

         # If it doesn't support cache, pass the test
         if not hasattr(a__    ,					"""use_cache"""  ):
           return

         _lowerCAmelCase						:	Dict        =					model_class(a__  ).to(a__  )
         if "use_cache" not in inputs:
           _lowerCAmelCase						:	Union[str, Any]        =					True
         _lowerCAmelCase						:	int        =					model(**a__  )

         # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
         if "past_key_values" not in outputs:
           return

         _lowerCAmelCase						:	int        =					(
             getattr(a__    ,					"""decoder_layers"""    ,					a__  )
             or getattr(a__    ,					"""num_decoder_layers"""    ,					a__  )
             or config.num_hidden_layers
         )
         _lowerCAmelCase						:	Dict        =					getattr(a__    ,					"""num_kv_heads"""    ,					config.num_attention_heads  )
         _lowerCAmelCase						:	Optional[Any]        =					getattr(a__    ,					"""d_model"""    ,					config.hidden_size  )
         _lowerCAmelCase						:	List[Any]        =					embed_dim // num_attention_heads

         _lowerCAmelCase						:	int        =					outputs["""past_key_values"""]
         self.assertEqual(len(a__  )    ,					a__  )

         _lowerCAmelCase						:	Tuple        =					inputs["""input_ids"""].shape
         for i in range(a__  ):
           if config.new_decoder_architecture:
             _lowerCAmelCase						:	Any        =					config.num_attention_heads
           elif config.multi_query:
             _lowerCAmelCase						:	Dict        =					1
           self.assertEqual(len(past_kv[0]  )    ,					2  )  # K V for the decoder = 2
           self.assertEqual(
               past_kv[i][0].shape    ,					(batch_size, num_attention_heads, seq_length, per_head_embed_dim)  )
           self.assertEqual(
               past_kv[i][1].shape    ,					(batch_size, num_attention_heads, seq_length, per_head_embed_dim)  )







@require_torch
class 						__A							(			unittest.TestCase		):






     @slow
     def    __A  (		self  ):
       _lowerCAmelCase						:	Any        =					AutoTokenizer.from_pretrained("""Rocketknight1/falcon-rw-1b"""  )
       _lowerCAmelCase						:	Dict        =					FalconForCausalLM.from_pretrained("""Rocketknight1/falcon-rw-1b"""  )
       model.eval()
       model.to(a__  )
       _lowerCAmelCase						:	List[Any]        =					tokenizer("""My favorite food is"""    ,					return_tensors="""pt"""  ).to(a__  )

       _lowerCAmelCase						:	Dict        =					(
           """My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday."""
       )

       _lowerCAmelCase						:	List[Any]        =					model.generate(**a__    ,					do_sample=a__    ,					max_new_tokens=19  )
       _lowerCAmelCase						:	int        =					tokenizer.batch_decode(a__  )[0]

       self.assertEqual(a__    ,					a__  )






     @slow
     def    __A  (		self  ):
       # The big models are way too big for the CI, so we use tiny random models that resemble their
       # architectures but with much smaller and fewer layers
       for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
         _lowerCAmelCase						:	List[Any]        =					AutoTokenizer.from_pretrained(a__  )
         _lowerCAmelCase						:	Any        =					FalconForCausalLM.from_pretrained(a__  )
         model.eval()
         model.to(a__  )
         _lowerCAmelCase						:	Union[str, Any]        =					tokenizer("""My favorite food is"""    ,					return_tensors="""pt"""  ).to(a__  )

         # We just test that these run without errors - the models are randomly initialized
         # and so the actual text outputs will be garbage
         model.generate(**a__    ,					do_sample=a__    ,					max_new_tokens=4  )
         model.generate(**a__    ,					do_sample=a__    ,					max_new_tokens=4  )
         model.generate(**a__    ,					num_beams=2    ,					max_new_tokens=4  )






     @slow
     def    __A  (		self  ):
       # The big models are way too big for the CI, so we use tiny random models that resemble their
       # architectures but with much smaller and fewer layers
       with torch.no_grad():
         for repo in [
             "Rocketknight1/falcon-rw-1b",
             "Rocketknight1/tiny-random-falcon-7b",
             "Rocketknight1/tiny-random-falcon-40b",
         ]:
           _lowerCAmelCase						:	int        =					AutoTokenizer.from_pretrained(a__  )
           _lowerCAmelCase						:	Union[str, Any]        =					FalconForCausalLM.from_pretrained(a__  )
           model.eval()
           model.to(device=a__  )
           _lowerCAmelCase						:	int        =					tokenizer("""My favorite food is"""    ,					return_tensors="""pt"""  ).to(a__  )

           # Test results are the same with and without cache
           _lowerCAmelCase						:	str        =					model.generate(**a__    ,					do_sample=a__    ,					max_new_tokens=20    ,					use_cache=a__  )
           _lowerCAmelCase						:	str        =					model.generate(**a__    ,					do_sample=a__    ,					max_new_tokens=20    ,					use_cache=a__  )
           self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0  )





"""simple docstring"""




import os
from pathlib import Path

import numpy as np
import pytest
from pack_dataset import pack_data_dir
from parameterized import parameterized
from save_len_file import save_len_file
from torch.utils.data import DataLoader

from transformers import AutoTokenizer
from transformers.models.mbart.modeling_mbart import shift_tokens_right
from transformers.testing_utils import TestCasePlus, slow
from utils import FAIRSEQ_AVAILABLE, DistributedSortishSampler, LegacySeqaSeqDataset, SeqaSeqDataset


_a				:		Optional[int]	   =       'bert-base-cased'
_a				:		Optional[Any]	   =       'google/pegasus-xsum'
_a				:		Union[str, Any]	   =       [' Sam ate lunch today.', 'Sams lunch ingredients.']
_a				:		int	   =       ['A very interesting story about what I ate for lunch.', 'Avocado, celery, turkey, coffee']
_a				:		Union[str, Any]	   =       'patrickvonplaten/t5-tiny-random'
_a				:		Tuple	   =       'sshleifer/bart-tiny-random'
_a				:		str	   =       'sshleifer/tiny-mbart'
_a				:		Optional[int]	   =       'sshleifer/tiny-marian-en-de'





def 							SCREAMING_SNAKE_CASE					(      _lowerCamelCase						:	Path      ,_lowerCamelCase						:	list     )	->				str:
  _lowerCAmelCase						:	List[Any]        =					"""\n""".join(_lowerCamelCase     )
  Path(_lowerCamelCase     ).open("""w"""     ).writelines(_lowerCamelCase     )





def 							SCREAMING_SNAKE_CASE					(      _lowerCamelCase						:	Optional[int]     )	->				Union[str, Any]:
  for split in ["train", "val", "test"]:
    _dump_articles(os.path.join(_lowerCamelCase      ,f"{split}.source"     )      ,_lowerCamelCase     )
    _dump_articles(os.path.join(_lowerCamelCase      ,f"{split}.target"     )      ,_lowerCamelCase     )
  return tmp_dir






class 						__A							(			SCREAMING_SNAKE_CASE_		):






     @parameterized.expand(
         [
             MBART_TINY,
             MARIAN_TINY,
             T5_TINY,
             BART_TINY,
             PEGASUS_XSUM,
         ]    ,					)
     @slow
     def    __A  (		self    ,					a__  ):
       _lowerCAmelCase						:	str        =					AutoTokenizer.from_pretrained(a__  )
       _lowerCAmelCase						:	Any        =					make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()  )
       _lowerCAmelCase						:	Union[str, Any]        =					max(len(tokenizer.encode(a__  )  ) for a in ARTICLES  )
       _lowerCAmelCase						:	Optional[int]        =					max(len(tokenizer.encode(a__  )  ) for a in SUMMARIES  )
       _lowerCAmelCase						:	str        =					4
       _lowerCAmelCase						:	Optional[int]        =					8
       assert max_len_target > max_src_len  # Will be truncated
       assert max_len_source > max_src_len  # Will be truncated
       _lowerCAmelCase      ,     _lowerCAmelCase						:	Optional[int]        =					"""ro_RO""", """de_DE"""  # ignored for all but mbart, but never causes error.
       _lowerCAmelCase						:	Optional[int]        =					SeqaSeqDataset(
           a__    ,					data_dir=a__    ,					type_path="""train"""    ,					max_source_length=a__    ,					max_target_length=a__    ,					src_lang=a__    ,					tgt_lang=a__    ,					)
       _lowerCAmelCase						:	int        =					DataLoader(a__    ,					batch_size=2    ,					collate_fn=train_dataset.collate_fn  )
       for batch in dataloader:
         assert isinstance(a__    ,					a__  )
         assert batch["attention_mask"].shape == batch["input_ids"].shape
         # show that articles were trimmed.
         assert batch["input_ids"].shape[1] == max_src_len
         # show that targets are the same len
         assert batch["labels"].shape[1] == max_tgt_len
         if tok_name != MBART_TINY:
           continue
         # check language codes in correct place
         _lowerCAmelCase						:	Any        =					shift_tokens_right(batch["""labels"""]    ,					tokenizer.pad_token_id  )
         assert batch["decoder_input_ids"][0, 0].item() == tokenizer.lang_code_to_id[tgt_lang]
         assert batch["decoder_input_ids"][0, -1].item() == tokenizer.eos_token_id
         assert batch["input_ids"][0, -2].item() == tokenizer.eos_token_id
         assert batch["input_ids"][0, -1].item() == tokenizer.lang_code_to_id[src_lang]

         break  # No need to test every batch






     @parameterized.expand([BART_TINY, BERT_BASE_CASED]  )
     def    __A  (		self    ,					a__  ):
       _lowerCAmelCase						:	Any        =					AutoTokenizer.from_pretrained(a__  )
       _lowerCAmelCase						:	Union[str, Any]        =					make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()  )
       _lowerCAmelCase						:	Optional[int]        =					max(len(tokenizer.encode(a__  )  ) for a in ARTICLES  )
       _lowerCAmelCase						:	Any        =					max(len(tokenizer.encode(a__  )  ) for a in SUMMARIES  )
       _lowerCAmelCase						:	int        =					4
       _lowerCAmelCase						:	List[str]        =					LegacySeqaSeqDataset(
           a__    ,					data_dir=a__    ,					type_path="""train"""    ,					max_source_length=20    ,					max_target_length=a__    ,					)
       _lowerCAmelCase						:	List[Any]        =					DataLoader(a__    ,					batch_size=2    ,					collate_fn=train_dataset.collate_fn  )
       for batch in dataloader:
         assert batch["attention_mask"].shape == batch["input_ids"].shape
         # show that articles were trimmed.
         assert batch["input_ids"].shape[1] == max_len_source
         assert 20 >= batch["input_ids"].shape[1]  # trimmed significantly
         # show that targets were truncated
         assert batch["labels"].shape[1] == trunc_target  # Truncated
         assert max_len_target > trunc_target  # Truncated
         break  # No need to test every batch






     def    __A  (		self  ):
       _lowerCAmelCase						:	Any        =					AutoTokenizer.from_pretrained("""facebook/mbart-large-cc25"""  )

       _lowerCAmelCase						:	List[Any]        =					Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()  )  )
       _lowerCAmelCase						:	List[Any]        =					tmp_dir.joinpath("""train.source"""  ).open().readlines()
       _lowerCAmelCase						:	str        =					Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()  )  )
       pack_data_dir(a__    ,					a__    ,					128    ,					a__  )
       _lowerCAmelCase						:	List[Any]        =					{x.name for x in tmp_dir.iterdir()}
       _lowerCAmelCase						:	Tuple        =					{x.name for x in save_dir.iterdir()}
       _lowerCAmelCase						:	Union[str, Any]        =					save_dir.joinpath("""train.source"""  ).open().readlines()
       # orig: [' Sam ate lunch today.\n', 'Sams lunch ingredients.']
       # desired_packed: [' Sam ate lunch today.\n Sams lunch ingredients.']
       assert len(a__  ) < len(a__  )
       assert len(a__  ) == 1
       assert len(packed_examples[0]  ) == sum(len(a__  ) for x in orig_examples  )
       assert orig_paths == new_paths






     @pytest.mark.skipif(not FAIRSEQ_AVAILABLE    ,					reason="""This test requires fairseq"""  )
     def    __A  (		self  ):
       if not FAIRSEQ_AVAILABLE:
         return
       _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase						:	Tuple        =					self._get_dataset(max_len=64  )
       _lowerCAmelCase						:	Optional[int]        =					64
       _lowerCAmelCase						:	str        =					ds.make_dynamic_sampler(a__    ,					required_batch_size_multiple=a__  )
       _lowerCAmelCase						:	int        =					[len(a__  ) for x in batch_sampler]
       assert len(set(a__  )  ) > 1  # it's not dynamic batch size if every batch is the same length
       assert sum(a__  ) == len(a__  )  # no dropped or added examples
       _lowerCAmelCase						:	List[str]        =					DataLoader(a__    ,					batch_sampler=a__    ,					collate_fn=ds.collate_fn    ,					num_workers=2  )
       _lowerCAmelCase						:	List[Any]        =					[]
       _lowerCAmelCase						:	Optional[int]        =					[]
       for batch in data_loader:
         _lowerCAmelCase						:	int        =					batch["""input_ids"""].shape
         _lowerCAmelCase						:	Union[str, Any]        =					src_shape[0]
         assert bs % required_batch_size_multiple == 0 or bs < required_batch_size_multiple
         _lowerCAmelCase						:	List[str]        =					np.product(batch["""input_ids"""].shape  )
         num_src_per_batch.append(a__  )
         if num_src_tokens > (max_tokens * 1.1):
           failures.append(a__  )
       assert num_src_per_batch[0] == max(a__  )
       if failures:
         raise AssertionError(F"too many tokens in {len(a__  )} batches"  )






     def    __A  (		self  ):
       _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase						:	Dict        =					self._get_dataset(max_len=512  )
       _lowerCAmelCase						:	int        =					2
       _lowerCAmelCase						:	List[str]        =					ds.make_sortish_sampler(a__    ,					shuffle=a__  )

       _lowerCAmelCase						:	Dict        =					DataLoader(a__    ,					batch_size=a__    ,					collate_fn=ds.collate_fn    ,					num_workers=2  )
       _lowerCAmelCase						:	int        =					DataLoader(a__    ,					batch_size=a__    ,					collate_fn=ds.collate_fn    ,					num_workers=2    ,					sampler=a__  )

       _lowerCAmelCase						:	int        =					tokenizer.pad_token_id

       def count_pad_tokens(a__    ,					a__="input_ids"  ):
         return [batch[k].eq(a__  ).sum().item() for batch in data_loader]

       assert sum(count_pad_tokens(a__    ,					k="""labels"""  )  ) < sum(count_pad_tokens(a__    ,					k="""labels"""  )  )
       assert sum(count_pad_tokens(a__  )  ) < sum(count_pad_tokens(a__  )  )
       assert len(a__  ) == len(a__  )






     def    __A  (		self    ,					a__=1000    ,					a__=128  ):
       if os.getenv("""USE_REAL_DATA"""    ,					a__  ):
         _lowerCAmelCase						:	List[str]        =					"""examples/seq2seq/wmt_en_ro"""
         _lowerCAmelCase						:	str        =					max_len * 2 * 64
         if not Path(a__  ).joinpath("""train.len"""  ).exists():
           save_len_file(a__    ,					a__  )
       else:
         _lowerCAmelCase						:	List[str]        =					"""examples/seq2seq/test_data/wmt_en_ro"""
         _lowerCAmelCase						:	Dict        =					max_len * 4
         save_len_file(a__    ,					a__  )

       _lowerCAmelCase						:	Optional[Any]        =					AutoTokenizer.from_pretrained(a__  )
       _lowerCAmelCase						:	Dict        =					SeqaSeqDataset(
           a__    ,					data_dir=a__    ,					type_path="""train"""    ,					max_source_length=a__    ,					max_target_length=a__    ,					n_obs=a__    ,					)
       return ds, max_tokens, tokenizer






     def    __A  (		self  ):
       _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase						:	Union[str, Any]        =					self._get_dataset()
       _lowerCAmelCase						:	Any        =					set(DistributedSortishSampler(a__    ,					256    ,					num_replicas=2    ,					rank=0    ,					add_extra_examples=a__  )  )
       _lowerCAmelCase						:	Optional[int]        =					set(DistributedSortishSampler(a__    ,					256    ,					num_replicas=2    ,					rank=1    ,					add_extra_examples=a__  )  )
       assert idsa.intersection(a__  ) == set()






     @parameterized.expand(
         [
             MBART_TINY,
             MARIAN_TINY,
             T5_TINY,
             BART_TINY,
             PEGASUS_XSUM,
         ]    ,					)
     def    __A  (		self    ,					a__  ):
       _lowerCAmelCase						:	int        =					AutoTokenizer.from_pretrained(a__    ,					use_fast=a__  )
       if tok_name == MBART_TINY:
         _lowerCAmelCase						:	Dict        =					SeqaSeqDataset(
             a__    ,					data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()  )    ,					type_path="""train"""    ,					max_source_length=4    ,					max_target_length=8    ,					src_lang="""EN"""    ,					tgt_lang="""FR"""    ,					)
         _lowerCAmelCase						:	Optional[Any]        =					train_dataset.dataset_kwargs
         assert "src_lang" in kwargs and "tgt_lang" in kwargs
       else:
         _lowerCAmelCase						:	List[Any]        =					SeqaSeqDataset(
             a__    ,					data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()  )    ,					type_path="""train"""    ,					max_source_length=4    ,					max_target_length=8    ,					)
         _lowerCAmelCase						:	Tuple        =					train_dataset.dataset_kwargs
         assert "add_prefix_space" not in kwargs if tok_name != BART_TINY else "add_prefix_space" in kwargs
         assert len(a__  ) == 1 if tok_name == BART_TINY else len(a__  ) == 0




