






import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch

def   _lowerCAmelCase   (						A__	):  # picklable for multiprocessing
      return x.sum()

def   _lowerCAmelCase   (						A__	):  # picklable for multiprocessing
      return i + 1



@dataclass
class       UpperCAmelCase__:

   '''simple docstring'''




   A    :  int
   A    :  str



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   def 			UpperCAmelCase		(						self  :			int)	->				Optional[Any]:




         """simple docstring"""

         lowercase__     =  {}
         lowercase__     =  []
         lowercase__     =  1
         lowercase__     =  [1, 2]
         lowercase__     =  {'''a''': 1, '''b''': 2}
         lowercase__     =  {'''a''': [1, 2], '''b''': [3, 4]}
         lowercase__     =  {'''a''': {'''1''': 1}, '''b''': 2}
         lowercase__     =  {'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
         lowercase__     =  {}
         lowercase__     =  []
         lowercase__     =  2
         lowercase__     =  [2, 3]
         lowercase__     =  {'''a''': 2, '''b''': 3}
         lowercase__     =  {'''a''': [2, 3], '''b''': [4, 5]}
         lowercase__     =  {'''a''': {'''1''': 2}, '''b''': 3}
         lowercase__     =  {'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__)			,				UpperCAmelCase__)

         lowercase__     =  2
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				num_proc=UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				num_proc=UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				num_proc=UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				num_proc=UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				num_proc=UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				num_proc=UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				num_proc=UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				num_proc=UpperCAmelCase__)			,				UpperCAmelCase__)

         lowercase__     =  {'''a''': np.eye(2), '''b''': np.zeros(3), '''c''': np.ones(2)}
         lowercase__     =  {'''a''': 2, '''b''': 0, '''c''': 2}
         lowercase__     =  {
             '''a''': np.eye(2).astype(UpperCAmelCase__),
             '''b''': np.zeros(3).astype(UpperCAmelCase__),
             '''c''': np.ones(2).astype(UpperCAmelCase__),
         }
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				map_numpy=UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(
             {k: v.tolist() for k, v in map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				map_numpy=UpperCAmelCase__).items()}			,				{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}			,				)
         self.assertEqual(map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				map_numpy=UpperCAmelCase__			,				num_proc=UpperCAmelCase__)			,				UpperCAmelCase__)
         self.assertEqual(
             {k: v.tolist() for k, v in map_nested(UpperCAmelCase__			,				UpperCAmelCase__			,				map_numpy=UpperCAmelCase__			,				num_proc=UpperCAmelCase__).items()}			,				{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}			,				)
         with self.assertRaises(UpperCAmelCase__):  # can't pickle a local lambda
               map_nested(lambda lowerCAmelCase: x + 1			,				UpperCAmelCase__			,				num_proc=UpperCAmelCase__)


   def 			UpperCAmelCase		(						self  :			Any)	->				Optional[int]:




         """simple docstring"""

         lowercase__     =  {'''a''': 1, '''b''': 2}
         lowercase__     =  {'''a''': 3, '''b''': 4}
         lowercase__     =  {'''a''': 5, '''b''': 6}
         lowercase__     =  sorted([('a', (1, 3, 5)), ('b', (2, 4, 6))])
         self.assertEqual(sorted(zip_dict(UpperCAmelCase__			,				UpperCAmelCase__			,				UpperCAmelCase__))			,				UpperCAmelCase__)



   def 			UpperCAmelCase		(						self  :			List[str])	->				Dict:




         """simple docstring"""

         class       UpperCAmelCase__:

            '''simple docstring'''




            A    :  str       				=    'bar'
         lowercase__     =  Foo()
         self.assertEqual(foo.my_attr			,				'bar')
         with temporary_assignment(UpperCAmelCase__			,				'my_attr'			,				'BAR'):
               self.assertEqual(foo.my_attr			,				'BAR')
         self.assertEqual(foo.my_attr			,				'bar')

@pytest.mark.parametrize(
    'iterable_length, num_proc, expected_num_proc'						,  [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]						,  )
def   _lowerCAmelCase   (						A__						,  A__						,  A__	):






      with patch('datasets.utils.py_utils._single_map_nested'	) as mock_single_map_nested, patch(
          'datasets.parallel.parallel.Pool'	) as mock_multiprocessing_pool:
            lowercase__     =  {F'''{i}''': i for i in range(A__	)}
            lowercase__     =  map_nested(lambda A__	: x + 10						,  A__						,  num_proc=A__						,  parallel_min_length=16	)
            if expected_num_proc == 1:
                  assert mock_single_map_nested.called
                  assert not mock_multiprocessing_pool.called
            else:
                  assert not mock_single_map_nested.called
                  assert mock_multiprocessing_pool.called
                  assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   @require_tf
   def 			UpperCAmelCase		(						self  :			List[str])	->				Union[str, Any]:




         """simple docstring"""

         import tensorflow as tf
         from tensorflow.keras import layers

         lowercase__     =  layers.Dense(2)

         def gen_random_output():
               lowercase__     =  tf.random.uniform((1, 3))
               return model(UpperCAmelCase__).numpy()

         with temp_seed(42			,				set_tensorflow=UpperCAmelCase__):
               lowercase__     =  gen_random_output()
         with temp_seed(42			,				set_tensorflow=UpperCAmelCase__):
               lowercase__     =  gen_random_output()
         lowercase__     =  gen_random_output()

         np.testing.assert_equal(UpperCAmelCase__			,				UpperCAmelCase__)
         self.assertGreater(np.abs(outa - outa).sum()			,				0)


   @require_torch
   def 			UpperCAmelCase		(						self  :			Optional[int])	->				Tuple:




         """simple docstring"""

         import torch

         def gen_random_output():
               lowercase__     =  torch.nn.Linear(3			,				2)
               lowercase__     =  torch.rand(1			,				3)
               return model(UpperCAmelCase__).detach().numpy()

         with temp_seed(42			,				set_pytorch=UpperCAmelCase__):
               lowercase__     =  gen_random_output()
         with temp_seed(42			,				set_pytorch=UpperCAmelCase__):
               lowercase__     =  gen_random_output()
         lowercase__     =  gen_random_output()

         np.testing.assert_equal(UpperCAmelCase__			,				UpperCAmelCase__)
         self.assertGreater(np.abs(outa - outa).sum()			,				0)



   def 			UpperCAmelCase		(						self  :			Dict)	->				Dict:




         """simple docstring"""

         def gen_random_output():
               return np.random.rand(1			,				3)

         with temp_seed(42):
               lowercase__     =  gen_random_output()
         with temp_seed(42):
               lowercase__     =  gen_random_output()
         lowercase__     =  gen_random_output()

         np.testing.assert_equal(UpperCAmelCase__			,				UpperCAmelCase__)
         self.assertGreater(np.abs(outa - outa).sum()			,				0)

@pytest.mark.parametrize('input_data'						,  [{}]	)
def   _lowerCAmelCase   (						A__	):
      lowercase__     =  NestedDataStructure(A__	).data
      assert output_data == input_data

@pytest.mark.parametrize(
    'data, expected_output'						,  [
        ({}, []),
        ([], []),
        ('foo', ['foo']),
        (['foo', 'bar'], ['foo', 'bar']),
        ([['foo', 'bar']], ['foo', 'bar']),
        ([[['foo'], ['bar']]], ['foo', 'bar']),
        ([[['foo'], 'bar']], ['foo', 'bar']),
        ({'a': 1, 'b': 2}, [1, 2]),
        ({'a': [1, 2], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[1, 2]], 'b': [[3, 4]]}, [1, 2, 3, 4]),
        ({'a': [[1, 2]], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [[[3], [4]]]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [[3, 4]]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [3, [4]]}, [1, 2, 3, 4]),
        ({'a': {'1': 1}, 'b': 2}, [1, 2]),
        ({'a': {'1': [1]}, 'b': 2}, [1, 2]),
        ({'a': {'1': [1]}, 'b': [2]}, [1, 2]),
    ]						,  )
def   _lowerCAmelCase   (						A__						,  A__	):
      lowercase__     =  NestedDataStructure(A__	).flatten()
      assert output == expected_output

def   _lowerCAmelCase   (						):
      lowercase__     =  A(x=1						,  y='foobar'	)
      lowercase__     =  {'''x''': 1, '''y''': '''foobar'''}
      assert asdict(A__	) == expected_output

      lowercase__     =  {'''a''': {'''b''': A(x=10						,  y='foo'	)}, '''c''': [A(x=20						,  y='bar'	)]}
      lowercase__     =  {'''a''': {'''b''': {'''x''': 10, '''y''': '''foo'''}}, '''c''': [{'''x''': 20, '''y''': '''bar'''}]}
      assert asdict(A__	) == expected_output

      with pytest.raises(A__	):
            asdict([1, A(x=10						,  y='foo'	)]	)

def   _lowerCAmelCase   (						A__	):
      return text.split()

def   _lowerCAmelCase   (						A__	):
      yield (time.time(), content)
      time.sleep(2	)
      yield (time.time(), content)

def   _lowerCAmelCase   (						):






      with Pool(2	) as pool:
            lowercase__     =  list(iflatmap_unordered(A__						,  _split_text						,  kwargs_iterable=[{'text': 'hello there'}] * 10	)	)
            assert out.count('hello'	) == 10
            assert out.count('there'	) == 10
            assert len(A__	) == 20

      # check multiprocess from pathos (uses dill for pickling)
      with multiprocess.Pool(2	) as pool:
            lowercase__     =  list(iflatmap_unordered(A__						,  _split_text						,  kwargs_iterable=[{'text': 'hello there'}] * 10	)	)
            assert out.count('hello'	) == 10
            assert out.count('there'	) == 10
            assert len(A__	) == 20

      # check that we get items as fast as possible
      with Pool(2	) as pool:
            lowercase__     =  []
            for yield_time, content in iflatmap_unordered(
                A__						,  _aseconds_generator_of_aitems_with_timing						,  kwargs_iterable=[{'content': 'a'}, {'content': 'b'}]	):
                  assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
                  out.append(A__	)
            assert out.count('a'	) == 2
            assert out.count('b'	) == 2
            assert len(A__	) == 4





import json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_bart import BartTokenizer


a__      :				List[Any]   	=			logging.get_logger(__name__)


a__      :				Optional[int]   	=			{"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}

# See all BART models at https://huggingface.co/models?filter=bart
a__      :				List[Any]   	=			{
    "vocab_file": {
        "facebook/bart-base": "https://huggingface.co/facebook/bart-base/resolve/main/vocab.json",
        "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/vocab.json",
        "facebook/bart-large-mnli": "https://huggingface.co/facebook/bart-large-mnli/resolve/main/vocab.json",
        "facebook/bart-large-cnn": "https://huggingface.co/facebook/bart-large-cnn/resolve/main/vocab.json",
        "facebook/bart-large-xsum": "https://huggingface.co/facebook/bart-large-xsum/resolve/main/vocab.json",
        "yjernite/bart_eli5": "https://huggingface.co/yjernite/bart_eli5/resolve/main/vocab.json",
    },
    "merges_file": {
        "facebook/bart-base": "https://huggingface.co/facebook/bart-base/resolve/main/merges.txt",
        "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/merges.txt",
        "facebook/bart-large-mnli": "https://huggingface.co/facebook/bart-large-mnli/resolve/main/merges.txt",
        "facebook/bart-large-cnn": "https://huggingface.co/facebook/bart-large-cnn/resolve/main/merges.txt",
        "facebook/bart-large-xsum": "https://huggingface.co/facebook/bart-large-xsum/resolve/main/merges.txt",
        "yjernite/bart_eli5": "https://huggingface.co/yjernite/bart_eli5/resolve/main/merges.txt",
    },
    "tokenizer_file": {
        "facebook/bart-base": "https://huggingface.co/facebook/bart-base/resolve/main/tokenizer.json",
        "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/tokenizer.json",
        "facebook/bart-large-mnli": "https://huggingface.co/facebook/bart-large-mnli/resolve/main/tokenizer.json",
        "facebook/bart-large-cnn": "https://huggingface.co/facebook/bart-large-cnn/resolve/main/tokenizer.json",
        "facebook/bart-large-xsum": "https://huggingface.co/facebook/bart-large-xsum/resolve/main/tokenizer.json",
        "yjernite/bart_eli5": "https://huggingface.co/yjernite/bart_eli5/resolve/main/tokenizer.json",
    },
}

a__      :				int   	=			{
    "facebook/bart-base": 10_24,
    "facebook/bart-large": 10_24,
    "facebook/bart-large-mnli": 10_24,
    "facebook/bart-large-cnn": 10_24,
    "facebook/bart-large-xsum": 10_24,
    "yjernite/bart_eli5": 10_24,
}



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   A    :  Optional[Any]       				=    VOCAB_FILES_NAMES
   A    :  Dict       				=    PRETRAINED_VOCAB_FILES_MAP
   A    :  str       				=    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   A    :  int       				=    ["input_ids", "attention_mask"]
   A    :  Any       				=    BartTokenizer


   def __init__(						self  :			List[Any]			,				lowerCAmelCase  :			Any=None			,				lowerCAmelCase  :			Optional[int]=None			,				lowerCAmelCase  :			List[Any]=None			,				lowerCAmelCase  :			str="replace"			,				lowerCAmelCase  :			str="<s>"			,				lowerCAmelCase  :			int="</s>"			,				lowerCAmelCase  :			Optional[int]="</s>"			,				lowerCAmelCase  :			Union[str, Any]="<s>"			,				lowerCAmelCase  :			str="<unk>"			,				lowerCAmelCase  :			int="<pad>"			,				lowerCAmelCase  :			int="<mask>"			,				lowerCAmelCase  :			Dict=False			,				lowerCAmelCase  :			List[Any]=True			,				**lowerCAmelCase  :			Optional[Any]			,				)	->				Optional[Any]:




         """simple docstring"""

         super().__init__(
             lowerCAmelCase			,				lowerCAmelCase			,				tokenizer_file=lowerCAmelCase			,				errors=lowerCAmelCase			,				bos_token=lowerCAmelCase			,				eos_token=lowerCAmelCase			,				sep_token=lowerCAmelCase			,				cls_token=lowerCAmelCase			,				unk_token=lowerCAmelCase			,				pad_token=lowerCAmelCase			,				mask_token=lowerCAmelCase			,				add_prefix_space=lowerCAmelCase			,				trim_offsets=lowerCAmelCase			,				**lowerCAmelCase			,				)

         lowercase__     =  json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__())
         if pre_tok_state.get('add_prefix_space'			,				lowerCAmelCase) != add_prefix_space:
               lowercase__     =  getattr(lowerCAmelCase			,				pre_tok_state.pop('type'))
               lowercase__     =  add_prefix_space
               lowercase__     =  pre_tok_class(**lowerCAmelCase)

         lowercase__     =  add_prefix_space

         # the pre_tokenizer is already updated in the GPT2TokenizerFast `__init__`
         lowercase__     =  'post_processor'
         lowercase__     =  getattr(self.backend_tokenizer			,				lowerCAmelCase			,				lowerCAmelCase)
         if tokenizer_component_instance:
               lowercase__     =  json.loads(tokenizer_component_instance.__getstate__())

               # The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
               if "sep" in state:
                     lowercase__     =  tuple(state['sep'])
               if "cls" in state:
                     lowercase__     =  tuple(state['cls'])

               lowercase__     =  False

               if state.get('add_prefix_space'			,				lowerCAmelCase) != add_prefix_space:
                     lowercase__     =  add_prefix_space
                     lowercase__     =  True

               if state.get('trim_offsets'			,				lowerCAmelCase) != trim_offsets:
                     lowercase__     =  trim_offsets
                     lowercase__     =  True

               if changes_to_apply:
                     lowercase__     =  getattr(lowerCAmelCase			,				state.pop('type'))
                     lowercase__     =  component_class(**lowerCAmelCase)
                     setattr(self.backend_tokenizer			,				lowerCAmelCase			,				lowerCAmelCase)


   @property
   def 			UpperCAmelCase		(						self  :			Union[str, Any])	->				str:




         """simple docstring"""

         if self._mask_token is None:
               if self.verbose:
                     logger.error('Using mask_token, but it is not set yet.')
               return None
         return str(self._mask_token)


   @mask_token.setter
   def 			UpperCAmelCase		(						self  :			Tuple			,				lowerCAmelCase  :			int)	->				Optional[int]:




         """simple docstring"""

         lowercase__     =  AddedToken(lowerCAmelCase			,				lstrip=lowerCAmelCase			,				rstrip=lowerCAmelCase) if isinstance(lowerCAmelCase			,				lowerCAmelCase) else value
         lowercase__     =  value


   def 			UpperCAmelCase		(						self  :			List[str]			,				*lowerCAmelCase  :			Tuple			,				**lowerCAmelCase  :			Optional[int])	->				BatchEncoding:




         """simple docstring"""

         lowercase__     =  kwargs.get('is_split_into_words'			,				lowerCAmelCase)

         if is_split_into_words and not self.add_prefix_space:
               raise ValueError(
                   f'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
                   'to use it with pretokenized inputs.')

         return super()._batch_encode_plus(*lowerCAmelCase			,				**lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			str			,				*lowerCAmelCase  :			Tuple			,				**lowerCAmelCase  :			str)	->				BatchEncoding:




         """simple docstring"""

         lowercase__     =  kwargs.get('is_split_into_words'			,				lowerCAmelCase)

         if is_split_into_words and not self.add_prefix_space:
               raise ValueError(
                   f'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
                   'to use it with pretokenized inputs.')

         return super()._encode_plus(*lowerCAmelCase			,				**lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			Dict			,				lowerCAmelCase  :			str			,				lowerCAmelCase  :			Optional[str] = None)	->				Tuple[str]:




         """simple docstring"""

         lowercase__     =  self._tokenizer.model.save(lowerCAmelCase			,				name=lowerCAmelCase)
         return tuple(lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			Any			,				lowerCAmelCase  :			str			,				lowerCAmelCase  :			Optional[int]=None)	->				Tuple:




         """simple docstring"""

         lowercase__     =  [self.bos_token_id] + token_ids_a + [self.eos_token_id]
         if token_ids_a is None:
               return output

         return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]



   def 			UpperCAmelCase		(						self  :			Union[str, Any]			,				lowerCAmelCase  :			List[int]			,				lowerCAmelCase  :			Optional[List[int]] = None)	->				List[int]:




         """simple docstring"""

         lowercase__     =  [self.sep_token_id]
         lowercase__     =  [self.cls_token_id]

         if token_ids_a is None:
               return len(cls + token_ids_a + sep) * [0]
         return len(cls + token_ids_a + sep + sep + token_ids_a + sep) * [0]
