


"""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


__lowercase    					=      get_tests_dir("""fixtures/test_sentencepiece.model""")


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

__lowercase    					=      50003
__lowercase    					=      50002






@require_sentencepiece
@require_tokenizers
class 						_A  (  _a    ,unittest.TestCase  ):







    """simple docstring"""




    UpperCAmelCase       :							int	=       PLBartTokenizer
    UpperCAmelCase       :							List[Any]	=       None
    UpperCAmelCase       :							str	=       False






    def 	__snake_case							(					self				:    int):
      super().setUp()

      # We have a SentencePiece fixture for testing
      a						:		List[str]      							=			PLBartTokenizer(__UpperCAmelCase				, language_codes="base"				, keep_accents=__UpperCAmelCase)
      tokenizer.save_pretrained(self.tmpdirname)






    def 	__snake_case							(					self				:    List[str]):
      a						:		List[Any]      							=			PLBartTokenizer(__UpperCAmelCase				, language_codes="base"				, keep_accents=__UpperCAmelCase)

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

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

      a						:		Union[str, Any]      							=			tokenizer.tokenize("I was born in 92000, and this is falsé.")
      self.assertListEqual(
          __UpperCAmelCase				, [
              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",
              "é",
              ".",
          ]				, )
      a						:		Tuple      							=			tokenizer.convert_tokens_to_ids(__UpperCAmelCase)
      self.assertListEqual(
          __UpperCAmelCase				, [
              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]
          ]				, )

      a						:		Union[str, Any]      							=			tokenizer.convert_ids_to_tokens(__UpperCAmelCase)
      self.assertListEqual(
          __UpperCAmelCase				, [
              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>",
              ".",
          ]				, )

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

      self.assertListEqual(__UpperCAmelCase				, ["__java__", "__python__", "__en_XX__", "<mask>"])

      a						:		Union[str, Any]      							=			"java.lang.Exception, python.lang.Exception, javascript, php, ruby, go"
      a						:		Tuple      							=			tokenizer(__UpperCAmelCase).input_ids
      self.assertEqual(
          tokenizer.decode(__UpperCAmelCase				, skip_special_tokens=__UpperCAmelCase				, clean_up_tokenization_spaces=__UpperCAmelCase)				, __UpperCAmelCase				, )






    def 	__snake_case							(					self				:    int):
      a						:		Any      							=			PLBartTokenizer(__UpperCAmelCase				, language_codes="multi"				, keep_accents=__UpperCAmelCase)

      a						:		Optional[int]      							=			tokenizer.tokenize("This is a test")
      self.assertListEqual(__UpperCAmelCase				, ["▁This", "▁is", "▁a", "▁t", "est"])

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

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

      a						:		Dict      							=			tokenizer.convert_ids_to_tokens(__UpperCAmelCase)
      self.assertListEqual(
          __UpperCAmelCase				, [
              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>",
              ".",
          ]				, )
      a						:		Any      							=			tokenizer.vocab_size
      a						:		Tuple      							=			[tokenizer.convert_ids_to_tokens(__UpperCAmelCase) for x in range(end - 7				, __UpperCAmelCase)]

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






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







    """simple docstring"""




    UpperCAmelCase       :							int	=       """uclanlp/plbart-python-en_XX"""
    UpperCAmelCase       :							Dict	=       [
        """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       :							Optional[Any]	=       [
        """Returns the maximum value of a b c.""",
        """Sums the values of a b c.""",
    ]
    UpperCAmelCase       :							Dict	=       [
        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 	__snake_case							(					cls				:    Optional[Any]):
      a						:		PLBartTokenizer      							=			PLBartTokenizer.from_pretrained(
          cls.checkpoint_name				, language_codes="base"				, src_lang="python"				, tgt_lang="en_XX")
      a						:		Union[str, Any]      							=			1
      return cls






    def 	__snake_case							(					self				:    Any):
      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 	__snake_case							(					self				:    List[Any]):
      a						:		Any      							=			self.tokenizer.batch_encode_plus(self.src_text).input_ids[0]
      self.assertListEqual(self.expected_src_tokens				, __UpperCAmelCase)






    def 	__snake_case							(					self				:    Tuple):
      self.assertIn(__UpperCAmelCase				, self.tokenizer.all_special_ids)
      a						:		Tuple      							=			[EN_CODE, 9037, 33442, 57, 752, 153, 14, 56, 18, 9, 2]
      a						:		str      							=			self.tokenizer.decode(__UpperCAmelCase				, skip_special_tokens=__UpperCAmelCase)
      a						:		List[Any]      							=			self.tokenizer.decode(generated_ids[1:]				, skip_special_tokens=__UpperCAmelCase)
      self.assertEqual(__UpperCAmelCase				, __UpperCAmelCase)
      self.assertNotIn(self.tokenizer.eos_token				, __UpperCAmelCase)






    def 	__snake_case							(					self				:    Any):
      a						:		Dict      							=			["def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])" * 20]
      self.assertIsInstance(src_text[0]				, __UpperCAmelCase)
      a						:		List[str]      							=			10
      a						:		Union[str, Any]      							=			self.tokenizer(__UpperCAmelCase				, max_length=__UpperCAmelCase				, truncation=__UpperCAmelCase).input_ids[0]
      self.assertEqual(ids[-2]				, 2)
      self.assertEqual(ids[-1]				, __UpperCAmelCase)
      self.assertEqual(len(__UpperCAmelCase)				, __UpperCAmelCase)






    def 	__snake_case							(					self				:    int):
      self.assertListEqual(self.tokenizer.convert_tokens_to_ids(["<mask>", "__java__"])				, [50004, 50001])






    def 	__snake_case							(					self				:    Optional[int]):
      a						:		Union[str, Any]      							=			tempfile.mkdtemp()
      a						:		Dict      							=			self.tokenizer.fairseq_tokens_to_ids
      self.tokenizer.save_pretrained(__UpperCAmelCase)
      a						:		Optional[int]      							=			PLBartTokenizer.from_pretrained(__UpperCAmelCase)
      self.assertDictEqual(new_tok.fairseq_tokens_to_ids				, __UpperCAmelCase)






    @require_torch
    def 	__snake_case							(					self				:    Tuple):
      a						:		Optional[int]      							=			self.tokenizer(self.src_text				, text_target=self.tgt_text				, padding=__UpperCAmelCase				, return_tensors="pt")
      a						:		Optional[int]      							=			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]				, __UpperCAmelCase)
      self.assertEqual(batch.decoder_input_ids[1][-1]				, 2)
      self.assertEqual(batch.labels[1][-2:].tolist()				, [2, EN_CODE])






    @require_torch
    def 	__snake_case							(					self				:    Optional[Any]):
      a						:		Union[str, Any]      							=			self.tokenizer(
          self.src_text				, text_target=self.tgt_text				, padding=__UpperCAmelCase				, truncation=__UpperCAmelCase				, max_length=len(self.expected_src_tokens)				, return_tensors="pt"				, )
      a						:		Tuple      							=			shift_tokens_right(batch["labels"]				, self.tokenizer.pad_token_id)

      self.assertIsInstance(__UpperCAmelCase				, __UpperCAmelCase)

      self.assertEqual((2, 26)				, batch.input_ids.shape)
      self.assertEqual((2, 26)				, batch.attention_mask.shape)
      a						:		List[str]      							=			batch.input_ids.tolist()[0]
      self.assertListEqual(self.expected_src_tokens				, __UpperCAmelCase)
      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 	__snake_case							(					self				:    str):
      a						:		Tuple      							=			self.tokenizer(self.src_text				, padding=__UpperCAmelCase				, truncation=__UpperCAmelCase				, max_length=3				, return_tensors="pt")
      a						:		Optional[Any]      							=			self.tokenizer(
          text_target=self.tgt_text				, padding=__UpperCAmelCase				, truncation=__UpperCAmelCase				, max_length=10				, return_tensors="pt")
      a						:		Dict      							=			targets["input_ids"]
      a						:		str      							=			shift_tokens_right(__UpperCAmelCase				, 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 	__snake_case							(					self				:    Dict):
      a						:		Optional[int]      							=			self.tokenizer._build_translation_inputs(
          "A test"				, return_tensors="pt"				, src_lang="en_XX"				, tgt_lang="java")

      self.assertEqual(
          nested_simplify(__UpperCAmelCase)				, {
              # 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 itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


__lowercase    					=      datasets.utils.logging.get_logger(__name__)






@dataclass
class 						_A  (  datasets.BuilderConfig  ):







    """simple docstring"""




    UpperCAmelCase       :							int	=       1_0_0_0_0
    UpperCAmelCase       :							Optional[List[str]]	=       None
    UpperCAmelCase       :							Optional[datasets.Features]	=       None






class 						_A  (  datasets.ArrowBasedBuilder  ):







    """simple docstring"""




    UpperCAmelCase       :							str	=       ParquetConfig






    def 	__snake_case							(					self				:    Tuple):
      return datasets.DatasetInfo(features=self.config.features)






    def 	__snake_case							(					self				:    List[Any]				, __UpperCAmelCase				:    str):
      if not self.config.data_files:
        raise ValueError(f'''At least one data file must be specified, but got data_files={self.config.data_files}''')
      a						:		str      							=			dl_manager.download_and_extract(self.config.data_files)
      if isinstance(__UpperCAmelCase				, (str, list, tuple)):
        a						:		Dict      							=			data_files
        if isinstance(__UpperCAmelCase				, __UpperCAmelCase):
          a						:		str      							=			[files]
        # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
        a						:		List[Any]      							=			[dl_manager.iter_files(__UpperCAmelCase) for file in files]
        return [datasets.SplitGenerator(name=datasets.Split.TRAIN				, gen_kwargs={"files": files})]
      a						:		Dict      							=			[]
      for split_name, files in data_files.items():
        if isinstance(__UpperCAmelCase				, __UpperCAmelCase):
          a						:		Optional[int]      							=			[files]
        # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
        a						:		Tuple      							=			[dl_manager.iter_files(__UpperCAmelCase) for file in files]
        # Infer features is they are stoed in the arrow schema
        if self.info.features is None:
          for file in itertools.chain.from_iterable(__UpperCAmelCase):
            with open(__UpperCAmelCase				, "rb") as f:
              a						:		Tuple      							=			datasets.Features.from_arrow_schema(pq.read_schema(__UpperCAmelCase))
            break
        splits.append(datasets.SplitGenerator(name=__UpperCAmelCase				, gen_kwargs={"files": files}))
      return splits






    def 	__snake_case							(					self				:    List[str]				, __UpperCAmelCase				:    pa.Table):
      if self.info.features is not None:
        # more expensive cast to support nested features with keys in a different order
        # allows str <-> int/float or str to Audio for example
        a						:		Optional[int]      							=			table_cast(__UpperCAmelCase				, self.info.features.arrow_schema)
      return pa_table






    def 	__snake_case							(					self				:    Union[str, Any]				, __UpperCAmelCase				:    int):
      a						:		Tuple      							=			self.info.features.arrow_schema if self.info.features is not None else None
      if self.info.features is not None and self.config.columns is not None:
        if sorted(field.name for field in schema) != sorted(self.config.columns):
          raise ValueError(
              f'''Tried to load parquet data with columns \'{self.config.columns}\' with mismatching features \'{self.info.features}\'''')
      for file_idx, file in enumerate(itertools.chain.from_iterable(__UpperCAmelCase)):
        with open(__UpperCAmelCase				, "rb") as f:
          a						:		Tuple      							=			pq.ParquetFile(__UpperCAmelCase)
          try:
            for batch_idx, record_batch in enumerate(
                parquet_file.iter_batches(batch_size=self.config.batch_size				, columns=self.config.columns)):
              a						:		Optional[Any]      							=			pa.Table.from_batches([record_batch])
              # Uncomment for debugging (will print the Arrow table size and elements)
              # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
              # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
              yield f'''{file_idx}_{batch_idx}''', self._cast_table(__UpperCAmelCase)
          except ValueError as e:
            logger.error(f'''Failed to read file \'{file}\' with error {type(__UpperCAmelCase)}: {e}''')
            raise






