

import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_electra import ElectraTokenizer


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

UpperCamelCase__  :      Any			=    {
    "vocab_file": {
        "google/electra-small-generator": (
            "https://huggingface.co/google/electra-small-generator/resolve/main/vocab.txt"
        ),
        "google/electra-base-generator": "https://huggingface.co/google/electra-base-generator/resolve/main/vocab.txt",
        "google/electra-large-generator": (
            "https://huggingface.co/google/electra-large-generator/resolve/main/vocab.txt"
        ),
        "google/electra-small-discriminator": (
            "https://huggingface.co/google/electra-small-discriminator/resolve/main/vocab.txt"
        ),
        "google/electra-base-discriminator": (
            "https://huggingface.co/google/electra-base-discriminator/resolve/main/vocab.txt"
        ),
        "google/electra-large-discriminator": (
            "https://huggingface.co/google/electra-large-discriminator/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "google/electra-small-generator": (
            "https://huggingface.co/google/electra-small-generator/resolve/main/tokenizer.json"
        ),
        "google/electra-base-generator": (
            "https://huggingface.co/google/electra-base-generator/resolve/main/tokenizer.json"
        ),
        "google/electra-large-generator": (
            "https://huggingface.co/google/electra-large-generator/resolve/main/tokenizer.json"
        ),
        "google/electra-small-discriminator": (
            "https://huggingface.co/google/electra-small-discriminator/resolve/main/tokenizer.json"
        ),
        "google/electra-base-discriminator": (
            "https://huggingface.co/google/electra-base-discriminator/resolve/main/tokenizer.json"
        ),
        "google/electra-large-discriminator": (
            "https://huggingface.co/google/electra-large-discriminator/resolve/main/tokenizer.json"
        ),
    },
}

UpperCamelCase__  :      List[Any]			=    {
    "google/electra-small-generator": 512,
    "google/electra-base-generator": 512,
    "google/electra-large-generator": 512,
    "google/electra-small-discriminator": 512,
    "google/electra-base-discriminator": 512,
    "google/electra-large-discriminator": 512,
}

UpperCamelCase__  :      List[str]			=    {
    "google/electra-small-generator": {"do_lower_case": True},
    "google/electra-base-generator": {"do_lower_case": True},
    "google/electra-large-generator": {"do_lower_case": True},
    "google/electra-small-discriminator": {"do_lower_case": True},
    "google/electra-base-discriminator": {"do_lower_case": True},
    "google/electra-large-discriminator": {"do_lower_case": True},
}



class       __snake_case  (    lowercase__       ):
		__lowerCAmelCase     :	Optional[Any]						 =  VOCAB_FILES_NAMES
		__lowerCAmelCase     :	Dict						 =  PRETRAINED_VOCAB_FILES_MAP
		__lowerCAmelCase     :	Union[str, Any]						 =  PRETRAINED_INIT_CONFIGURATION
		__lowerCAmelCase     :	Tuple						 =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		__lowerCAmelCase     :	Union[str, Any]						 =  ElectraTokenizer

		def __init__(					self						,		_A=None						,		_A=None						,		_A=True						,		_A="[UNK]"						,		_A="[SEP]"						,		_A="[PAD]"						,		_A="[CLS]"						,		_A="[MASK]"						,		_A=True						,		_A=None						,		**_A						,		):
						super().__init__(
						    __lowercase						,		tokenizer_file=__lowercase						,		do_lower_case=__lowercase						,		unk_token=__lowercase						,		sep_token=__lowercase						,		pad_token=__lowercase						,		cls_token=__lowercase						,		mask_token=__lowercase						,		tokenize_chinese_chars=__lowercase						,		strip_accents=__lowercase						,		**__lowercase						,		)

						SCREAMING_SNAKE_CASE_									=						json.loads(self.backend_tokenizer.normalizer.__getstate__())
						if (
						    normalizer_state.get('lowercase'						,		__lowercase) != do_lower_case
						    or normalizer_state.get('strip_accents'						,		__lowercase) != strip_accents
						    or normalizer_state.get('handle_chinese_chars'						,		__lowercase) != tokenize_chinese_chars
						):
										SCREAMING_SNAKE_CASE_									=						getattr(__lowercase						,		normalizer_state.pop('type'))
										SCREAMING_SNAKE_CASE_									=						do_lower_case
										SCREAMING_SNAKE_CASE_									=						strip_accents
										SCREAMING_SNAKE_CASE_									=						tokenize_chinese_chars
										SCREAMING_SNAKE_CASE_									=						normalizer_class(**__lowercase)

						SCREAMING_SNAKE_CASE_									=						do_lower_case

		def 						lowerCAmelCase__     (					self						,		_A						,		_A=None):
						SCREAMING_SNAKE_CASE_									=						[self.cls_token_id] + token_ids_a + [self.sep_token_id]

						if token_ids_a:
										output += token_ids_a + [self.sep_token_id]

						return output

		def 						lowerCAmelCase__     (					self						,		_A						,		_A = None):
						SCREAMING_SNAKE_CASE_									=						[self.sep_token_id]
						SCREAMING_SNAKE_CASE_									=						[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) * [0] + len(token_ids_a + sep) * [1]


		def 						lowerCAmelCase__     (					self						,		_A						,		_A = None):
						SCREAMING_SNAKE_CASE_									=						self._tokenizer.model.save(__lowercase						,		name=__lowercase)
						return tuple(__lowercase)





import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


UpperCamelCase__  :      Optional[int]			=    datasets.utils.logging.get_logger(__name__)



@dataclass
class       __snake_case  (    datasets.BuilderConfig       ):
  __lowerCAmelCase     :	Optional[datasets.Features]						 =  None
  __lowerCAmelCase     :	str						 =  "utf-8"
  __lowerCAmelCase     :	Optional[str]						 =  None
  __lowerCAmelCase     :	Optional[str]						 =  None
  __lowerCAmelCase     :	bool						 =  True  # deprecated
  __lowerCAmelCase     :	Optional[int]						 =  None  # deprecated
  __lowerCAmelCase     :	int						 =  10 << 20  # 10MB
  __lowerCAmelCase     :	Optional[bool]						 =  None






class       __snake_case  (    datasets.ArrowBasedBuilder       ):
  __lowerCAmelCase     :	int						 =  JsonConfig

  def 						lowerCAmelCase__     (					self):
      if self.config.block_size is not None:
          logger.warning('The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead')
          SCREAMING_SNAKE_CASE_									=						self.config.block_size
      if self.config.use_threads is not True:
          logger.warning(
              'The JSON loader parameter `use_threads` is deprecated and doesn\'t have any effect anymore.')
      if self.config.newlines_in_values is not None:
          raise ValueError('The JSON loader parameter `newlines_in_values` is no longer supported')
      return datasets.DatasetInfo(features=self.config.features)

  def 						lowerCAmelCase__     (					self						,		_A):
      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}""")
      SCREAMING_SNAKE_CASE_									=						dl_manager.download_and_extract(self.config.data_files)
      if isinstance(_A						,		(str, list, tuple)):
          SCREAMING_SNAKE_CASE_									=						data_files
          if isinstance(_A						,		_A):
              SCREAMING_SNAKE_CASE_									=						[files]
          SCREAMING_SNAKE_CASE_									=						[dl_manager.iter_files(_A) for file in files]
          return [datasets.SplitGenerator(name=datasets.Split.TRAIN						,		gen_kwargs={'files': files})]
      SCREAMING_SNAKE_CASE_									=						[]
      for split_name, files in data_files.items():
          if isinstance(_A						,		_A):
              SCREAMING_SNAKE_CASE_									=						[files]
          SCREAMING_SNAKE_CASE_									=						[dl_manager.iter_files(_A) for file in files]
          splits.append(datasets.SplitGenerator(name=_A						,		gen_kwargs={'files': files}))
      return splits

  def 						lowerCAmelCase__     (					self						,		_A):
      if self.config.features is not None:
          # adding missing columns
          for column_name in set(self.config.features) - set(pa_table.column_names):
              SCREAMING_SNAKE_CASE_									=						self.config.features.arrow_schema.field(_A).type
              SCREAMING_SNAKE_CASE_									=						pa_table.append_column(_A						,		pa.array([None] * len(_A)						,		type=_A))
          # more expensive cast to support nested structures with keys in a different order
          # allows str <-> int/float or str to Audio for example
          SCREAMING_SNAKE_CASE_									=						table_cast(_A						,		self.config.features.arrow_schema)
      return pa_table


  def 						lowerCAmelCase__     (					self						,		_A):
      for file_idx, file in enumerate(itertools.chain.from_iterable(_A)):
          # If the file is one json object and if we need to look at the list of items in one specific field
          if self.config.field is not None:
              with open(_A						,		encoding=self.config.encoding						,		errors=self.config.encoding_errors) as f:
                  SCREAMING_SNAKE_CASE_									=						json.load(_A)

              # We keep only the field we are interested in
              SCREAMING_SNAKE_CASE_									=						dataset[self.config.field]

              # We accept two format: a list of dicts or a dict of lists
              if isinstance(_A						,		(list, tuple)):
                  SCREAMING_SNAKE_CASE_									=						set().union(*[row.keys() for row in dataset])
                  SCREAMING_SNAKE_CASE_									=						{col: [row.get(_A) for row in dataset] for col in keys}
              else:
                  SCREAMING_SNAKE_CASE_									=						dataset
              SCREAMING_SNAKE_CASE_									=						pa.Table.from_pydict(_A)
              yield file_idx, self._cast_table(_A)

          # If the file has one json object per line
          else:
              with open(_A						,		'rb') as f:
                  SCREAMING_SNAKE_CASE_									=						0
                  # Use block_size equal to the chunk size divided by 32 to leverage multithreading
                  # Set a default minimum value of 16kB if the chunk size is really small
                  SCREAMING_SNAKE_CASE_									=						max(self.config.chunksize // 32						,		16 << 10)
                  SCREAMING_SNAKE_CASE_									=						(
                      self.config.encoding_errors if self.config.encoding_errors is not None else 'strict'
                  )
                  while True:
                      SCREAMING_SNAKE_CASE_									=						f.read(self.config.chunksize)
                      if not batch:
                          break
                      # Finish current line
                      try:
                          batch += f.readline()
                      except (AttributeError, io.UnsupportedOperation):
                          batch += readline(_A)
                      # PyArrow only accepts utf-8 encoded bytes
                      if self.config.encoding != "utf-8":
                          SCREAMING_SNAKE_CASE_									=						batch.decode(self.config.encoding						,		errors=_A).encode('utf-8')
                      try:
                          while True:
                              try:
                                  SCREAMING_SNAKE_CASE_									=						paj.read_json(
                                      io.BytesIO(_A)						,		read_options=paj.ReadOptions(block_size=_A))
                                  break
                              except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
                                  if (
                                      isinstance(_A						,		pa.ArrowInvalid)
                                      and "straddling" not in str(_A)
                                      or block_size > len(_A)
                                  ):
                                      raise
                                  else:
                                      # Increase the block size in case it was too small.
                                      # The block size will be reset for the next file.
                                      logger.debug(
                                          f"""Batch of {len(_A)} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}.""")
                                      block_size *= 2
                      except pa.ArrowInvalid as e:
                          try:
                              with open(
                                  _A						,		encoding=self.config.encoding						,		errors=self.config.encoding_errors) as f:
                                  SCREAMING_SNAKE_CASE_									=						json.load(_A)
                          except json.JSONDecodeError:
                              logger.error(f"""Failed to read file '{file}' with error {type(_A)}: {e}""")
                              raise e
                          # If possible, parse the file as a list of json objects and exit the loop
                          if isinstance(_A						,		_A):  # list is the only sequence type supported in JSON
                              try:
                                  SCREAMING_SNAKE_CASE_									=						set().union(*[row.keys() for row in dataset])
                                  SCREAMING_SNAKE_CASE_									=						{col: [row.get(_A) for row in dataset] for col in keys}
                                  SCREAMING_SNAKE_CASE_									=						pa.Table.from_pydict(_A)
                              except (pa.ArrowInvalid, AttributeError) as e:
                                  logger.error(f"""Failed to read file '{file}' with error {type(_A)}: {e}""")
                                  raise ValueError(f"""Not able to read records in the JSON file at {file}.""") from None
                              yield file_idx, self._cast_table(_A)
                              break
                          else:
                              logger.error(f"""Failed to read file '{file}' with error {type(_A)}: {e}""")
                              raise ValueError(
                                  f"""Not able to read records in the JSON file at {file}. """
                                  f"""You should probably indicate the field of the JSON file containing your records. """
                                  f"""This JSON file contain the following fields: {str(list(dataset.keys()))}. """
                                  f"""Select the correct one and provide it as `field='XXX'` to the dataset loading method. """) from None
                        # 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 (file_idx, batch_idx), self._cast_table(_A)
                      batch_idx += 1


