






"""simple docstring"""






import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, PLBartTokenizer, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
)

from ...test_tokenization_common import TokenizerTesterMixin


UpperCAmelCase_     :				Tuple         =				get_tests_dir("""fixtures/test_sentencepiece.model""")


if is_torch_available():
  from transformers.models.plbart.modeling_plbart import shift_tokens_right

UpperCAmelCase_     :				Union[str, Any]         =				50003
UpperCAmelCase_     :				str         =				50002




@require_sentencepiece
@require_tokenizers
class 							lowerCAmelCase__       (    UpperCAmelCase__       ,	unittest.TestCase				):






      '''simple docstring'''

      __UpperCamelCase							=						PLBartTokenizer
      __UpperCamelCase							=						None
      __UpperCamelCase							=						False
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          super().setUp()

          # We have a SentencePiece fixture for testing
          SCREAMING_SNAKE_CASE_   :						Tuple							=	PLBartTokenizer(lowercase_		,			language_codes='''base'''		,			keep_accents=lowercase_)
          tokenizer.save_pretrained(self.tmpdirname)
      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	PLBartTokenizer(lowercase_		,			language_codes='''base'''		,			keep_accents=lowercase_)

          SCREAMING_SNAKE_CASE_   :						Dict							=	tokenizer.tokenize('''This is a test''')
          self.assertListEqual(lowercase_		,			['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est'''])

          self.assertListEqual(
              tokenizer.convert_tokens_to_ids(lowercase_)		,			[value + tokenizer.fairseq_offset for value in [285, 46, 10, 170, 382]]		,			)

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	tokenizer.tokenize('''I was born in 92000, and this is falsé.''')
          self.assertListEqual(
              lowercase_		,			[
                  SPIECE_UNDERLINE + '''I''',
                  SPIECE_UNDERLINE + '''was''',
                  SPIECE_UNDERLINE + '''b''',
                  '''or''',
                  '''n''',
                  SPIECE_UNDERLINE + '''in''',
                  SPIECE_UNDERLINE + '''''',
                  '''9''',
                  '''2''',
                  '''0''',
                  '''0''',
                  '''0''',
                  ''',''',
                  SPIECE_UNDERLINE + '''and''',
                  SPIECE_UNDERLINE + '''this''',
                  SPIECE_UNDERLINE + '''is''',
                  SPIECE_UNDERLINE + '''f''',
                  '''al''',
                  '''s''',
                  '''é''',
                  '''.''',
              ]		,			)
          SCREAMING_SNAKE_CASE_   :						Dict							=	tokenizer.convert_tokens_to_ids(lowercase_)
          self.assertListEqual(
              lowercase_		,			[
                  value + tokenizer.fairseq_offset
                  for value in [8, 21, 84, 55, 24, 19, 7, 2, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 2, 4]
              ]		,			)

          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	tokenizer.convert_ids_to_tokens(lowercase_)
          self.assertListEqual(
              lowercase_		,			[
                  SPIECE_UNDERLINE + '''I''',
                  SPIECE_UNDERLINE + '''was''',
                  SPIECE_UNDERLINE + '''b''',
                  '''or''',
                  '''n''',
                  SPIECE_UNDERLINE + '''in''',
                  SPIECE_UNDERLINE + '''''',
                  '''<unk>''',
                  '''2''',
                  '''0''',
                  '''0''',
                  '''0''',
                  ''',''',
                  SPIECE_UNDERLINE + '''and''',
                  SPIECE_UNDERLINE + '''this''',
                  SPIECE_UNDERLINE + '''is''',
                  SPIECE_UNDERLINE + '''f''',
                  '''al''',
                  '''s''',
                  '''<unk>''',
                  '''.''',
              ]		,			)

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	tokenizer.vocab_size
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	[tokenizer.convert_ids_to_tokens(lowercase_) for x in range(end - 4		,			lowercase_)]

          self.assertListEqual(lowercase_		,			['''__java__''', '''__python__''', '''__en_XX__''', '''<mask>'''])

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	'''java.lang.Exception, python.lang.Exception, javascript, php, ruby, go'''
          SCREAMING_SNAKE_CASE_   :						List[str]							=	tokenizer(lowercase_).input_ids
          self.assertEqual(
              tokenizer.decode(lowercase_		,			skip_special_tokens=lowercase_		,			clean_up_tokenization_spaces=lowercase_)		,			lowercase_		,			)



      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	PLBartTokenizer(lowercase_		,			language_codes='''multi'''		,			keep_accents=lowercase_)

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	tokenizer.tokenize('''This is a test''')
          self.assertListEqual(lowercase_		,			['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est'''])

          self.assertListEqual(
              tokenizer.convert_tokens_to_ids(lowercase_)		,			[value + tokenizer.fairseq_offset for value in [285, 46, 10, 170, 382]]		,			)

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	tokenizer.tokenize('''I was born in 92000, and this is falsé.''')
          self.assertListEqual(
              lowercase_		,			[
                  SPIECE_UNDERLINE + '''I''',
                  SPIECE_UNDERLINE + '''was''',
                  SPIECE_UNDERLINE + '''b''',
                  '''or''',
                  '''n''',
                  SPIECE_UNDERLINE + '''in''',
                  SPIECE_UNDERLINE + '''''',
                  '''9''',
                  '''2''',
                  '''0''',
                  '''0''',
                  '''0''',
                  ''',''',
                  SPIECE_UNDERLINE + '''and''',
                  SPIECE_UNDERLINE + '''this''',
                  SPIECE_UNDERLINE + '''is''',
                  SPIECE_UNDERLINE + '''f''',
                  '''al''',
                  '''s''',
                  '''é''',
                  '''.''',
              ]		,			)
          SCREAMING_SNAKE_CASE_   :						List[str]							=	tokenizer.convert_tokens_to_ids(lowercase_)
          self.assertListEqual(
              lowercase_		,			[
                  value + tokenizer.fairseq_offset
                  for value in [8, 21, 84, 55, 24, 19, 7, 2, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 2, 4]
              ]		,			)

          SCREAMING_SNAKE_CASE_   :						List[str]							=	tokenizer.convert_ids_to_tokens(lowercase_)
          self.assertListEqual(
              lowercase_		,			[
                  SPIECE_UNDERLINE + '''I''',
                  SPIECE_UNDERLINE + '''was''',
                  SPIECE_UNDERLINE + '''b''',
                  '''or''',
                  '''n''',
                  SPIECE_UNDERLINE + '''in''',
                  SPIECE_UNDERLINE + '''''',
                  '''<unk>''',
                  '''2''',
                  '''0''',
                  '''0''',
                  '''0''',
                  ''',''',
                  SPIECE_UNDERLINE + '''and''',
                  SPIECE_UNDERLINE + '''this''',
                  SPIECE_UNDERLINE + '''is''',
                  SPIECE_UNDERLINE + '''f''',
                  '''al''',
                  '''s''',
                  '''<unk>''',
                  '''.''',
              ]		,			)
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	tokenizer.vocab_size
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	[tokenizer.convert_ids_to_tokens(lowercase_) for x in range(end - 7		,			lowercase_)]

          self.assertListEqual(
              lowercase_		,			['''__java__''', '''__python__''', '''__en_XX__''', '''__javascript__''', '''__php__''', '''__ruby__''', '''__go__'''])
          SCREAMING_SNAKE_CASE_   :						Dict							=	'''java.lang.Exception, python.lang.Exception, javascript, php, ruby, go'''
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	tokenizer(lowercase_).input_ids
          self.assertEqual(
              tokenizer.decode(lowercase_		,			skip_special_tokens=lowercase_		,			clean_up_tokenization_spaces=lowercase_)		,			lowercase_		,			)




@require_torch
@require_sentencepiece
@require_tokenizers
class 							lowerCAmelCase__       (    unittest.TestCase				):






      '''simple docstring'''

      __UpperCamelCase							=						"uclanlp/plbart-python-en_XX"
      __UpperCamelCase							=						[
          "def maximum(a,b,c):NEW_LINE_INDENTreturn max([a,b,c])",
          "def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])",
      ]
      __UpperCamelCase							=						[
          "Returns the maximum value of a b c.",
          "Sums the values of a b c.",
      ]
      __UpperCamelCase							=						[
          1_3_4,
          5_4_5_2,
          3_3_4_6_0,
          3_3_4_4_1,
          3_3_4_6_3,
          3_3_4_6_5,
          3_3_4_6_3,
          3_3_4_4_9,
          9_8_8,
          2_0,
          3_3_4_5_6,
          1_9,
          3_3_4_5_6,
          7_7_1,
          3_9,
          4_2_5_8,
          8_8_9,
          3_3_1_8,
          3_3_4_4_1,
          3_3_4_6_3,
          3_3_4_6_5,
          3_3_4_6_3,
          3_3_4_4_9,
          2_4_7_1,
          2,
          PYTHON_CODE,
      ]
      @classmethod
      def _SCREAMING_SNAKE_CASE				(		cls  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						PLBartTokenizer							=	PLBartTokenizer.from_pretrained(
              cls.checkpoint_name		,			language_codes='''base'''		,			src_lang='''python'''		,			tgt_lang='''en_XX''')
          SCREAMING_SNAKE_CASE_   :						int							=	1
          return cls
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''__java__''']		,			50001)
          self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''__python__''']		,			50002)
          self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''__en_XX__''']		,			50003)
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						str							=	self.tokenizer.batch_encode_plus(self.src_text).input_ids[0]
          self.assertListEqual(self.expected_src_tokens		,			lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          self.assertIn(lowercase_		,			self.tokenizer.all_special_ids)
          SCREAMING_SNAKE_CASE_   :						str							=	[EN_CODE, 9037, 33442, 57, 752, 153, 14, 56, 18, 9, 2]
          SCREAMING_SNAKE_CASE_   :						Dict							=	self.tokenizer.decode(lowercase_		,			skip_special_tokens=lowercase_)
          SCREAMING_SNAKE_CASE_   :						Dict							=	self.tokenizer.decode(generated_ids[1:]		,			skip_special_tokens=lowercase_)
          self.assertEqual(lowercase_		,			lowercase_)
          self.assertNotIn(self.tokenizer.eos_token		,			lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	['''def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])''' * 20]
          self.assertIsInstance(src_text[0]		,			lowercase_)
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	10
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	self.tokenizer(lowercase_		,			max_length=lowercase_		,			truncation=lowercase_).input_ids[0]
          self.assertEqual(ids[-2]		,			2)
          self.assertEqual(ids[-1]		,			lowercase_)
          self.assertEqual(len(lowercase_)		,			lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['''<mask>''', '''__java__'''])		,			[50004, 50001])
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Dict							=	tempfile.mkdtemp()
          SCREAMING_SNAKE_CASE_   :						List[str]							=	self.tokenizer.fairseq_tokens_to_ids
          self.tokenizer.save_pretrained(lowercase_)
          SCREAMING_SNAKE_CASE_   :						int							=	PLBartTokenizer.from_pretrained(lowercase_)
          self.assertDictEqual(new_tok.fairseq_tokens_to_ids		,			lowercase_)
      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	self.tokenizer(self.src_text		,			text_target=self.tgt_text		,			padding=lowercase_		,			return_tensors='''pt''')
          SCREAMING_SNAKE_CASE_   :						List[str]							=	shift_tokens_right(batch['''labels''']		,			self.tokenizer.pad_token_id)

          # fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
          self.assertEqual(batch.input_ids[1][-2:].tolist()		,			[2, PYTHON_CODE])
          self.assertEqual(batch.decoder_input_ids[1][0]		,			lowercase_)
          self.assertEqual(batch.decoder_input_ids[1][-1]		,			2)
          self.assertEqual(batch.labels[1][-2:].tolist()		,			[2, EN_CODE])
      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	self.tokenizer(
              self.src_text		,			text_target=self.tgt_text		,			padding=lowercase_		,			truncation=lowercase_		,			max_length=len(self.expected_src_tokens)		,			return_tensors='''pt'''		,			)
          SCREAMING_SNAKE_CASE_   :						str							=	shift_tokens_right(batch['''labels''']		,			self.tokenizer.pad_token_id)

          self.assertIsInstance(lowercase_		,			lowercase_)

          self.assertEqual((2, 26)		,			batch.input_ids.shape)
          self.assertEqual((2, 26)		,			batch.attention_mask.shape)
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	batch.input_ids.tolist()[0]
          self.assertListEqual(self.expected_src_tokens		,			lowercase_)
          self.assertEqual(2		,			batch.decoder_input_ids[0, -1])  # EOS
          # Test that special tokens are reset
          self.assertEqual(self.tokenizer.prefix_tokens		,			[])
          self.assertEqual(self.tokenizer.suffix_tokens		,			[self.tokenizer.eos_token_id, PYTHON_CODE])
      def _SCREAMING_SNAKE_CASE				(		self  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[str]							=	self.tokenizer(self.src_text		,			padding=lowercase_		,			truncation=lowercase_		,			max_length=3		,			return_tensors='''pt''')
          SCREAMING_SNAKE_CASE_   :						Tuple							=	self.tokenizer(
              text_target=self.tgt_text		,			padding=lowercase_		,			truncation=lowercase_		,			max_length=10		,			return_tensors='''pt''')
          SCREAMING_SNAKE_CASE_   :						int							=	targets['''input_ids''']
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	shift_tokens_right(lowercase_		,			self.tokenizer.pad_token_id)

          self.assertEqual(batch.input_ids.shape[1]		,			3)
          self.assertEqual(batch.decoder_input_ids.shape[1]		,			10)



      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[str]							=	self.tokenizer._build_translation_inputs(
              '''A test'''		,			return_tensors='''pt'''		,			src_lang='''en_XX'''		,			tgt_lang='''java''')

          self.assertEqual(
              nested_simplify(lowercase_)		,			{
                  # A, test, EOS, en_XX
                  '''input_ids''': [[150, 242, 2, 50003]],
                  '''attention_mask''': [[1, 1, 1, 1]],
                  # java
                  '''forced_bos_token_id''': 50001,
              }		,			)










"""simple docstring"""






import multiprocessing
import time

from arguments import PretokenizationArguments
from datasets import load_dataset

from transformers import AutoTokenizer, HfArgumentParser



def       _A    (__a							)				-> Optional[int]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	{}
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	tokenizer(example['''content''']					,       truncation=__a							)['''input_ids''']
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	len(example['''content''']							) / len(output['''input_ids''']							)
    return output


UpperCAmelCase_     :				Tuple         =				HfArgumentParser(PretokenizationArguments)
UpperCAmelCase_     :				Tuple         =				parser.parse_args()
if args.num_workers is None:
  UpperCAmelCase_     :				Tuple         =				multiprocessing.cpu_count()
UpperCAmelCase_     :				str         =				AutoTokenizer.from_pretrained(args.tokenizer_dir)

UpperCAmelCase_     :				Tuple         =				time.time()
UpperCAmelCase_     :				List[Any]         =				load_dataset(args.dataset_name, split="""train""")
print(f'''Dataset loaded in {time.time()-t_start:.2f}s''')

UpperCAmelCase_     :				Dict         =				time.time()
UpperCAmelCase_     :				Dict         =				ds.map(
    tokenize,
    num_proc=args.num_workers,
    remove_columns=[
        """repo_name""",
        """path""",
        """copies""",
        """size""",
        """content""",
        """license""",
        """hash""",
        """line_mean""",
        """line_max""",
        """alpha_frac""",
        """autogenerated""",
    ],
)
print(f'''Dataset tokenized in {time.time()-t_start:.2f}s''')

UpperCAmelCase_     :				str         =				time.time()
ds.push_to_hub(args.tokenized_data_repo)





print(f'''Data pushed to the hub in {time.time()-t_start:.2f}s''')


