


'''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__    					=      get_tests_dir('''fixtures/test_sentencepiece.model''')


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

UpperCamelCase__    					=      5_0_0_0_3
UpperCamelCase__    					=      5_0_0_0_2




@require_sentencepiece
@require_tokenizers
class 			lowerCamelCase_   (	__a     ,       unittest.TestCase      ):
      lowerCAmelCase__		     =     PLBartTokenizer
      lowerCAmelCase__		     =     None
      lowerCAmelCase__		     =     False






      def lowercase_							(					self				:    int ):


          '''simple docstring'''


          super().setUp()

          # We have a SentencePiece fixture for testing
          UpperCAmelCase__			:		int													=				PLBartTokenizer(_A						,						language_codes='''base'''						,						keep_accents=_A )
          tokenizer.save_pretrained(self.tmpdirname )






      def lowercase_							(					self				:    Dict ):


          '''simple docstring'''


          UpperCAmelCase__			:		Tuple													=				PLBartTokenizer(_A						,						language_codes='''base'''						,						keep_accents=_A )

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

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

          UpperCAmelCase__			:		int													=				tokenizer.tokenize('''I was born in 92000, and this is falsé.''' )
          self.assertListEqual(
              _A						,						[
                  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''',
                  '''é''',
                  '''.''',
              ]						,						)
          UpperCAmelCase__			:		Optional[int]													=				tokenizer.convert_tokens_to_ids(_A )
          self.assertListEqual(
              _A						,						[
                  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]
              ]						,						)

          UpperCAmelCase__			:		Any													=				tokenizer.convert_ids_to_tokens(_A )
          self.assertListEqual(
              _A						,						[
                  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>''',
                  '''.''',
              ]						,						)

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

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

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






      def lowercase_							(					self				:    Any ):


          '''simple docstring'''


          UpperCAmelCase__			:		List[Any]													=				PLBartTokenizer(_A						,						language_codes='''multi'''						,						keep_accents=_A )

          UpperCAmelCase__			:		Tuple													=				tokenizer.tokenize('''This is a test''' )
          self.assertListEqual(_A						,						['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est'''] )

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

          UpperCAmelCase__			:		List[Any]													=				tokenizer.tokenize('''I was born in 92000, and this is falsé.''' )
          self.assertListEqual(
              _A						,						[
                  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''',
                  '''é''',
                  '''.''',
              ]						,						)
          UpperCAmelCase__			:		List[str]													=				tokenizer.convert_tokens_to_ids(_A )
          self.assertListEqual(
              _A						,						[
                  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]
              ]						,						)

          UpperCAmelCase__			:		Optional[int]													=				tokenizer.convert_ids_to_tokens(_A )
          self.assertListEqual(
              _A						,						[
                  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>''',
                  '''.''',
              ]						,						)
          UpperCAmelCase__			:		Dict													=				tokenizer.vocab_size
          UpperCAmelCase__			:		Optional[int]													=				[tokenizer.convert_ids_to_tokens(_A ) for x in range(end - 7						,						_A )]

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




@require_torch
@require_sentencepiece
@require_tokenizers
class 			lowerCamelCase_   (	unittest.TestCase      ):
      lowerCAmelCase__		     =     'uclanlp/plbart-python-en_XX'
      lowerCAmelCase__		     =     [
          'def maximum(a,b,c):NEW_LINE_INDENTreturn max([a,b,c])',
          'def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])',
      ]
      lowerCAmelCase__		     =     [
          'Returns the maximum value of a b c.',
          'Sums the values of a b c.',
      ]
      lowerCAmelCase__		     =     [
          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 lowercase_							(					cls				:    str ):


          '''simple docstring'''


          UpperCAmelCase__			:		PLBartTokenizer													=				PLBartTokenizer.from_pretrained(
              cls.checkpoint_name						,						language_codes='''base'''						,						src_lang='''python'''						,						tgt_lang='''en_XX''' )
          UpperCAmelCase__			:		Optional[Any]													=				1
          return cls






      def lowercase_							(					self				:    List[Any] ):


          '''simple docstring'''


          self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''__java__''']						,						50_001 )
          self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''__python__''']						,						50_002 )
          self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''__en_XX__''']						,						50_003 )






      def lowercase_							(					self				:    Union[str, Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		Union[str, Any]													=				self.tokenizer.batch_encode_plus(self.src_text ).input_ids[0]
          self.assertListEqual(self.expected_src_tokens						,						_A )






      def lowercase_							(					self				:    Optional[Any] ):


          '''simple docstring'''


          self.assertIn(_A						,						self.tokenizer.all_special_ids )
          UpperCAmelCase__			:		Optional[Any]													=				[EN_CODE, 9_037, 33_442, 57, 752, 153, 14, 56, 18, 9, 2]
          UpperCAmelCase__			:		Any													=				self.tokenizer.decode(_A						,						skip_special_tokens=_A )
          UpperCAmelCase__			:		Dict													=				self.tokenizer.decode(generated_ids[1:]						,						skip_special_tokens=_A )
          self.assertEqual(_A						,						_A )
          self.assertNotIn(self.tokenizer.eos_token						,						_A )






      def lowercase_							(					self				:    Optional[int] ):


          '''simple docstring'''


          UpperCAmelCase__			:		Tuple													=				['''def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])''' * 20]
          self.assertIsInstance(src_text[0]						,						_A )
          UpperCAmelCase__			:		List[str]													=				10
          UpperCAmelCase__			:		int													=				self.tokenizer(_A						,						max_length=_A						,						truncation=_A ).input_ids[0]
          self.assertEqual(ids[-2]						,						2 )
          self.assertEqual(ids[-1]						,						_A )
          self.assertEqual(len(_A )						,						_A )






      def lowercase_							(					self				:    List[Any] ):


          '''simple docstring'''


          self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['''<mask>''', '''__java__'''] )						,						[50_004, 50_001] )






      def lowercase_							(					self				:    Any ):


          '''simple docstring'''


          UpperCAmelCase__			:		Union[str, Any]													=				tempfile.mkdtemp()
          UpperCAmelCase__			:		Optional[Any]													=				self.tokenizer.fairseq_tokens_to_ids
          self.tokenizer.save_pretrained(_A )
          UpperCAmelCase__			:		Dict													=				PLBartTokenizer.from_pretrained(_A )
          self.assertDictEqual(new_tok.fairseq_tokens_to_ids						,						_A )






      @require_torch
      def lowercase_							(					self				:    Tuple ):


          '''simple docstring'''


          UpperCAmelCase__			:		int													=				self.tokenizer(self.src_text						,						text_target=self.tgt_text						,						padding=_A						,						return_tensors='''pt''' )
          UpperCAmelCase__			:		Any													=				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]						,						_A )
          self.assertEqual(batch.decoder_input_ids[1][-1]						,						2 )
          self.assertEqual(batch.labels[1][-2:].tolist()						,						[2, EN_CODE] )






      @require_torch
      def lowercase_							(					self				:    Union[str, Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		int													=				self.tokenizer(
              self.src_text						,						text_target=self.tgt_text						,						padding=_A						,						truncation=_A						,						max_length=len(self.expected_src_tokens )						,						return_tensors='''pt'''						,						)
          UpperCAmelCase__			:		List[str]													=				shift_tokens_right(batch['''labels''']						,						self.tokenizer.pad_token_id )

          self.assertIsInstance(_A						,						_A )

          self.assertEqual((2, 26)						,						batch.input_ids.shape )
          self.assertEqual((2, 26)						,						batch.attention_mask.shape )
          UpperCAmelCase__			:		List[str]													=				batch.input_ids.tolist()[0]
          self.assertListEqual(self.expected_src_tokens						,						_A )
          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 lowercase_							(					self				:    Dict ):


          '''simple docstring'''


          UpperCAmelCase__			:		Union[str, Any]													=				self.tokenizer(self.src_text						,						padding=_A						,						truncation=_A						,						max_length=3						,						return_tensors='''pt''' )
          UpperCAmelCase__			:		Tuple													=				self.tokenizer(
              text_target=self.tgt_text						,						padding=_A						,						truncation=_A						,						max_length=10						,						return_tensors='''pt''' )
          UpperCAmelCase__			:		List[Any]													=				targets['''input_ids''']
          UpperCAmelCase__			:		List[str]													=				shift_tokens_right(_A						,						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 lowercase_							(					self				:    int ):


          '''simple docstring'''


          UpperCAmelCase__			:		str													=				self.tokenizer._build_translation_inputs(
              '''A test'''						,						return_tensors='''pt'''						,						src_lang='''en_XX'''						,						tgt_lang='''java''' )

          self.assertEqual(
              nested_simplify(_A )						,						{
                  # A, test, EOS, en_XX
                  '''input_ids''': [[150, 242, 2, 50_003]],
                  '''attention_mask''': [[1, 1, 1, 1]],
                  # java
                  '''forced_bos_token_id''': 50_001,
              }						,						)




'''simple docstring'''
import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader




class 			lowerCamelCase_   (	__a      ):






      def __init__(					self				:    List[str]						,						_A				:    NestedDataStructureLike[PathLike]						,						_A				:    Optional[NamedSplit] = None						,						_A				:    Optional[Features] = None						,						_A				:    str = None						,						_A				:    bool = False						,						_A				:    bool = False						,						_A				:    Optional[str] = None						,						_A				:    Optional[int] = None						,						**_A				:    Union[str, Any]						,						):


          '''simple docstring'''


          super().__init__(
              _A						,						split=_A						,						features=_A						,						cache_dir=_A						,						keep_in_memory=_A						,						streaming=_A						,						num_proc=_A						,						**_A						,						)
          UpperCAmelCase__			:		List[str]													=				field
          UpperCAmelCase__			:		Optional[Any]													=				path_or_paths if isinstance(_A						,						_A ) else {self.split: path_or_paths}
          UpperCAmelCase__			:		Any													=				Json(
              cache_dir=_A						,						data_files=_A						,						features=_A						,						field=_A						,						**_A						,						)






      def lowercase_							(					self				:    Dict ):


          '''simple docstring'''


          if self.streaming:
              UpperCAmelCase__			:		int													=				self.builder.as_streaming_dataset(split=self.split )
          # Build regular (map-style) dataset
          else:
              UpperCAmelCase__			:		List[str]													=				None
              UpperCAmelCase__			:		Dict													=				None
              UpperCAmelCase__			:		Optional[int]													=				None
              UpperCAmelCase__			:		List[Any]													=				None

              self.builder.download_and_prepare(
                  download_config=_A						,						download_mode=_A						,						verification_mode=_A						,						base_path=_A						,						num_proc=self.num_proc						,						)
              UpperCAmelCase__			:		str													=				self.builder.as_dataset(
                  split=self.split						,						verification_mode=_A						,						in_memory=self.keep_in_memory )
          return dataset




class 			lowerCamelCase_   :






      def __init__(					self				:    Optional[Any]						,						_A				:    Dataset						,						_A				:    Union[PathLike, BinaryIO]						,						_A				:    Optional[int] = None						,						_A				:    Optional[int] = None						,						**_A				:    List[str]						,						):


          '''simple docstring'''


          if num_proc is not None and num_proc <= 0:
              raise ValueError(f"""num_proc {num_proc} must be an integer > 0.""" )

          UpperCAmelCase__			:		Dict													=				dataset
          UpperCAmelCase__			:		Any													=				path_or_buf
          UpperCAmelCase__			:		List[str]													=				batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
          UpperCAmelCase__			:		Tuple													=				num_proc
          UpperCAmelCase__			:		Any													=				'''utf-8'''
          UpperCAmelCase__			:		Optional[int]													=				to_json_kwargs






      def lowercase_							(					self				:    int ):


          '''simple docstring'''


          UpperCAmelCase__			:		List[str]													=				self.to_json_kwargs.pop('''path_or_buf'''						,						_A )
          UpperCAmelCase__			:		Optional[int]													=				self.to_json_kwargs.pop('''orient'''						,						'''records''' )
          UpperCAmelCase__			:		Tuple													=				self.to_json_kwargs.pop('''lines'''						,						True if orient == '''records''' else False )
          UpperCAmelCase__			:		str													=				self.to_json_kwargs.pop('''index'''						,						False if orient in ['''split''', '''table'''] else True )
          UpperCAmelCase__			:		Optional[Any]													=				self.to_json_kwargs.pop('''compression'''						,						_A )

          if compression not in [None, "infer", "gzip", "bz2", "xz"]:
              raise NotImplementedError(f"""`datasets` currently does not support {compression} compression""" )

          if isinstance(self.path_or_buf						,						(str, bytes, os.PathLike) ):
              with fsspec.open(self.path_or_buf						,						'''wb'''						,						compression=_A ) as buffer:
                  UpperCAmelCase__			:		Union[str, Any]													=				self._write(file_obj=_A						,						orient=_A						,						lines=_A						,						index=_A						,						**self.to_json_kwargs )
          else:
              if compression:
                  raise NotImplementedError(
                      f"""The compression parameter is not supported when writing to a buffer, but compression={compression}"""
                      ''' was passed. Please provide a local path instead.''' )
              UpperCAmelCase__			:		Union[str, Any]													=				self._write(
                  file_obj=self.path_or_buf						,						orient=_A						,						lines=_A						,						index=_A						,						**self.to_json_kwargs )
          return written






      def lowercase_							(					self				:    Optional[int]						,						_A				:    Optional[Any] ):


          '''simple docstring'''


          UpperCAmelCase__						,						UpperCAmelCase__						,						UpperCAmelCase__						,						UpperCAmelCase__						,						UpperCAmelCase__			:		Dict													=				args

          UpperCAmelCase__			:		Dict													=				query_table(
              table=self.dataset.data						,						key=slice(_A						,						offset + self.batch_size )						,						indices=self.dataset._indices						,						)
          UpperCAmelCase__			:		str													=				batch.to_pandas().to_json(
              path_or_buf=_A						,						orient=_A						,						lines=_A						,						index=_A						,						**_A )
          if not json_str.endswith('''\n''' ):
              json_str += "\n"
          return json_str.encode(self.encoding )






      def lowercase_							(					self				:    Union[str, Any]						,						_A				:    BinaryIO						,						_A				:    Optional[int]						,						_A				:    int						,						_A				:    Any						,						**_A				:    Any						,						):


          '''simple docstring'''


          UpperCAmelCase__			:		Optional[Any]													=				0

          if self.num_proc is None or self.num_proc == 1:
              for offset in logging.tqdm(
                  range(0						,						len(self.dataset )						,						self.batch_size )						,						unit='''ba'''						,						disable=not logging.is_progress_bar_enabled()						,						desc='''Creating json from Arrow format'''						,						):
                  UpperCAmelCase__			:		Optional[int]													=				self._batch_json((offset, orient, lines, index, to_json_kwargs) )
                  written += file_obj.write(_A )
          else:
              UpperCAmelCase__						,						UpperCAmelCase__			:		str													=				len(self.dataset ), self.batch_size
              with multiprocessing.Pool(self.num_proc ) as pool:
                  for json_str in logging.tqdm(
                      pool.imap(
                          self._batch_json						,						[(offset, orient, lines, index, to_json_kwargs) for offset in range(0						,						_A						,						_A )]						,						)						,						total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size						,						unit='''ba'''						,						disable=not logging.is_progress_bar_enabled()						,						desc='''Creating json from Arrow format'''						,						):
                      written += file_obj.write(_A )

          return written
