



import json
import os
from functools import lru_cache
from typing import List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


_UpperCAmelCase	       =	logging.get_logger(__name__)


_UpperCAmelCase	       =	{'''vocab_file''': '''vocab.json''', '''merges_file''': '''merges.txt'''}

_UpperCAmelCase	       =	{
    '''vocab_file''': {
        '''allenai/longformer-base-4096''': '''https://huggingface.co/allenai/longformer-base-4096/resolve/main/vocab.json''',
        '''allenai/longformer-large-4096''': (
            '''https://huggingface.co/allenai/longformer-large-4096/resolve/main/vocab.json'''
        ),
        '''allenai/longformer-large-4096-finetuned-triviaqa''': (
            '''https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/vocab.json'''
        ),
        '''allenai/longformer-base-4096-extra.pos.embd.only''': (
            '''https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/vocab.json'''
        ),
        '''allenai/longformer-large-4096-extra.pos.embd.only''': (
            '''https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/vocab.json'''
        ),
    },
    '''merges_file''': {
        '''allenai/longformer-base-4096''': '''https://huggingface.co/allenai/longformer-base-4096/resolve/main/merges.txt''',
        '''allenai/longformer-large-4096''': (
            '''https://huggingface.co/allenai/longformer-large-4096/resolve/main/merges.txt'''
        ),
        '''allenai/longformer-large-4096-finetuned-triviaqa''': (
            '''https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/merges.txt'''
        ),
        '''allenai/longformer-base-4096-extra.pos.embd.only''': (
            '''https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/merges.txt'''
        ),
        '''allenai/longformer-large-4096-extra.pos.embd.only''': (
            '''https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/merges.txt'''
        ),
    },
}

_UpperCAmelCase	       =	{
    '''allenai/longformer-base-4096''': 4_096,
    '''allenai/longformer-large-4096''': 4_096,
    '''allenai/longformer-large-4096-finetuned-triviaqa''': 4_096,
    '''allenai/longformer-base-4096-extra.pos.embd.only''': 4_096,
    '''allenai/longformer-large-4096-extra.pos.embd.only''': 4_096,
}
@lru_cache()
# Copied from transformers.models.roberta.tokenization_roberta.bytes_to_unicode
def 		UpperCamelCase     (       )     ->					Union[str, Any]:
   UpperCamelCase   :  Dict						     =     (
       list(range(ord('!'     )	,					ord('~'     ) + 1     )     ) + list(range(ord('¡'     )	,					ord('¬'     ) + 1     )     ) + list(range(ord('®'     )	,					ord('ÿ'     ) + 1     )     )
   )
   UpperCamelCase   :  Tuple						     =     bs[:]
   UpperCamelCase   :  List[Any]						     =     0
   for b in range(2**8     ):
      if b not in bs:
         bs.append(snake_case__     )
         cs.append(2**8 + n     )
         n += 1
   UpperCamelCase   :  Optional[Any]						     =     [chr(snake_case__     ) for n in cs]
   return dict(zip(snake_case__	,					snake_case__     )     )
def 		UpperCamelCase     (       snake_case__	:   Tuple     )     ->					Tuple:
   UpperCamelCase   :  int						     =     set()
   UpperCamelCase   :  Any						     =     word[0]
   for char in word[1:]:
      pairs.add((prev_char, char)     )
      UpperCamelCase   :  str						     =     char
   return pairs



class     lowerCAmelCase_		(  a__							):
       UpperCAmelCase__    : Optional[int]       					=  VOCAB_FILES_NAMES
       UpperCAmelCase__    : str       					=  PRETRAINED_VOCAB_FILES_MAP
       UpperCAmelCase__    : List[str]       					=  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       UpperCAmelCase__    : List[str]       					=  ["input_ids", "attention_mask"]
       def __init__(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_="replace",      SCREAMING_SNAKE_CASE_="<s>",      SCREAMING_SNAKE_CASE_="</s>",      SCREAMING_SNAKE_CASE_="</s>",      SCREAMING_SNAKE_CASE_="<s>",      SCREAMING_SNAKE_CASE_="<unk>",      SCREAMING_SNAKE_CASE_="<pad>",      SCREAMING_SNAKE_CASE_="<mask>",      SCREAMING_SNAKE_CASE_=False,      **SCREAMING_SNAKE_CASE_,      )					->		Dict:
          UpperCamelCase   :  int						     =     AddedToken(SCREAMING_SNAKE_CASE_,      lstrip=SCREAMING_SNAKE_CASE_,      rstrip=SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	) else bos_token
          UpperCamelCase   :  Dict						     =     AddedToken(SCREAMING_SNAKE_CASE_,      lstrip=SCREAMING_SNAKE_CASE_,      rstrip=SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	) else eos_token
          UpperCamelCase   :  Optional[int]						     =     AddedToken(SCREAMING_SNAKE_CASE_,      lstrip=SCREAMING_SNAKE_CASE_,      rstrip=SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	) else sep_token
          UpperCamelCase   :  str						     =     AddedToken(SCREAMING_SNAKE_CASE_,      lstrip=SCREAMING_SNAKE_CASE_,      rstrip=SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	) else cls_token
          UpperCamelCase   :  int						     =     AddedToken(SCREAMING_SNAKE_CASE_,      lstrip=SCREAMING_SNAKE_CASE_,      rstrip=SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	) else unk_token
          UpperCamelCase   :  Optional[int]						     =     AddedToken(SCREAMING_SNAKE_CASE_,      lstrip=SCREAMING_SNAKE_CASE_,      rstrip=SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	) else pad_token

          # Mask token behave like a normal word, i.e. include the space before it
          UpperCamelCase   :  Any						     =     AddedToken(SCREAMING_SNAKE_CASE_,      lstrip=SCREAMING_SNAKE_CASE_,      rstrip=SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	) else mask_token

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

          with open(SCREAMING_SNAKE_CASE_,      encoding='utf-8'	) as vocab_handle:
             UpperCamelCase   :  Union[str, Any]						     =     json.load(SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Dict						     =     {v: k for k, v in self.encoder.items()}
          UpperCamelCase   :  Tuple						     =     errors  # how to handle errors in decoding
          UpperCamelCase   :  Optional[int]						     =     bytes_to_unicode()
          UpperCamelCase   :  Any						     =     {v: k for k, v in self.byte_encoder.items()}
          with open(SCREAMING_SNAKE_CASE_,      encoding='utf-8'	) as merges_handle:
             UpperCamelCase   :  Optional[int]						     =     merges_handle.read().split('\n'	)[1:-1]
          UpperCamelCase   :  Any						     =     [tuple(merge.split()	) for merge in bpe_merges]
          UpperCamelCase   :  List[str]						     =     dict(zip(SCREAMING_SNAKE_CASE_,      range(len(SCREAMING_SNAKE_CASE_	)	)	)	)
          UpperCamelCase   :  int						     =     {}
          UpperCamelCase   :  str						     =     add_prefix_space

          # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
          UpperCamelCase   :  int						     =     re.compile(R'\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+'	)
       @property
       def  snake_case_	(      self	)					->		Optional[Any]:
          return len(self.encoder	)
       def  snake_case_	(      self	)					->		Any:
          return dict(self.encoder,      **self.added_tokens_encoder	)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		Tuple:
          if token in self.cache:
             return self.cache[token]
          UpperCamelCase   :  int						     =     tuple(SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Union[str, Any]						     =     get_pairs(SCREAMING_SNAKE_CASE_	)

          if not pairs:
             return token

          while True:
             UpperCamelCase   :  Dict						     =     min(SCREAMING_SNAKE_CASE_,      key=lambda SCREAMING_SNAKE_CASE_	: self.bpe_ranks.get(SCREAMING_SNAKE_CASE_,      float('inf'	)	)	)
             if bigram not in self.bpe_ranks:
                break
             UpperCamelCase   :  Any						     =     bigram
             UpperCamelCase   :  Dict						     =     []
             UpperCamelCase   :  Dict						     =     0
             while i < len(SCREAMING_SNAKE_CASE_	):
                try:
                   UpperCamelCase   :  Any						     =     word.index(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
                except ValueError:
                   new_word.extend(word[i:]	)
                   break
                else:
                   new_word.extend(word[i:j]	)
                   UpperCamelCase   :  int						     =     j

                if word[i] == first and i < len(SCREAMING_SNAKE_CASE_	) - 1 and word[i + 1] == second:
                   new_word.append(first + second	)
                   i += 2
                else:
                   new_word.append(word[i]	)
                   i += 1
             UpperCamelCase   :  Optional[Any]						     =     tuple(SCREAMING_SNAKE_CASE_	)
             UpperCamelCase   :  List[str]						     =     new_word
             if len(SCREAMING_SNAKE_CASE_	) == 1:
                break
             else:
                UpperCamelCase   :  Optional[Any]						     =     get_pairs(SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  int						     =     ' '.join(SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Optional[int]						     =     word
          return word
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		List[str]:
          UpperCamelCase   :  Dict						     =     []
          for token in re.findall(self.pat,      SCREAMING_SNAKE_CASE_	):
             UpperCamelCase   :  Union[str, Any]						     =     ''.join(
                 self.byte_encoder[b] for b in token.encode('utf-8'	)	)  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
             bpe_tokens.extend(bpe_token for bpe_token in self.bpe(SCREAMING_SNAKE_CASE_	).split(' '	)	)
          return bpe_tokens
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		Tuple:
          return self.encoder.get(SCREAMING_SNAKE_CASE_,      self.encoder.get(self.unk_token	)	)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		Optional[int]:
          return self.decoder.get(SCREAMING_SNAKE_CASE_	)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		Union[str, Any]:
          UpperCamelCase   :  int						     =     ''.join(SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Dict						     =     bytearray([self.byte_decoder[c] for c in text]	).decode('utf-8',      errors=self.errors	)
          return text
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = None	)					->		Tuple[str]:
          if not os.path.isdir(SCREAMING_SNAKE_CASE_	):
             logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""	)
             return
          UpperCamelCase   :  str						     =     os.path.join(
              SCREAMING_SNAKE_CASE_,      (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']	)
          UpperCamelCase   :  Any						     =     os.path.join(
              SCREAMING_SNAKE_CASE_,      (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file']	)

          with open(SCREAMING_SNAKE_CASE_,      'w',      encoding='utf-8'	) as f:
             f.write(json.dumps(self.encoder,      indent=2,      sort_keys=SCREAMING_SNAKE_CASE_,      ensure_ascii=SCREAMING_SNAKE_CASE_	) + '\n'	)

          UpperCamelCase   :  Dict						     =     0
          with open(SCREAMING_SNAKE_CASE_,      'w',      encoding='utf-8'	) as writer:
             writer.write('#version: 0.2\n'	)
             for bpe_tokens, token_index in sorted(self.bpe_ranks.items(),      key=lambda SCREAMING_SNAKE_CASE_	: kv[1]	):
                if index != token_index:
                   logger.warning(
                       F"""Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."""
                       ' Please check that the tokenizer is not corrupted!'	)
                   UpperCamelCase   :  Tuple						     =     token_index
                writer.write(' '.join(SCREAMING_SNAKE_CASE_	) + '\n'	)
                index += 1

          return vocab_file, merge_file
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = None	)					->		List[int]:
          if token_ids_a is None:
             return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
          UpperCamelCase   :  Any						     =     [self.cls_token_id]
          UpperCamelCase   :  int						     =     [self.sep_token_id]
          return cls + token_ids_a + sep + sep + token_ids_a + sep
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = None,      SCREAMING_SNAKE_CASE_ = False	)					->		List[int]:
          if already_has_special_tokens:
             return super().get_special_tokens_mask(
                 token_ids_a=SCREAMING_SNAKE_CASE_,      token_ids_a=SCREAMING_SNAKE_CASE_,      already_has_special_tokens=SCREAMING_SNAKE_CASE_	)

          if token_ids_a is None:
             return [1] + ([0] * len(SCREAMING_SNAKE_CASE_	)) + [1]
          return [1] + ([0] * len(SCREAMING_SNAKE_CASE_	)) + [1, 1] + ([0] * len(SCREAMING_SNAKE_CASE_	)) + [1]
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = None	)					->		List[int]:
          UpperCamelCase   :  List[Any]						     =     [self.sep_token_id]
          UpperCamelCase   :  Dict						     =     [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]






       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_=False,      **SCREAMING_SNAKE_CASE_	)					->		Optional[int]:
          UpperCamelCase   :  Tuple						     =     kwargs.pop('add_prefix_space',      self.add_prefix_space	)
          if (is_split_into_words or add_prefix_space) and (len(SCREAMING_SNAKE_CASE_	) > 0 and not text[0].isspace()):
             UpperCamelCase   :  Dict						     =     ' ' + text
          return (text, kwargs)








from typing import Optional

import pyspark

from .. import Features, NamedSplit
from ..download import DownloadMode
from ..packaged_modules.spark.spark import Spark
from .abc import AbstractDatasetReader



class     lowerCAmelCase_		(  a__							):
       def __init__(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = None,      SCREAMING_SNAKE_CASE_ = None,      SCREAMING_SNAKE_CASE_ = True,      SCREAMING_SNAKE_CASE_ = None,      SCREAMING_SNAKE_CASE_ = False,      SCREAMING_SNAKE_CASE_ = None,      SCREAMING_SNAKE_CASE_ = True,      SCREAMING_SNAKE_CASE_ = "arrow",      **SCREAMING_SNAKE_CASE_,      )					->		int:
          super().__init__(
              split=SCREAMING_SNAKE_CASE_,      features=SCREAMING_SNAKE_CASE_,      cache_dir=SCREAMING_SNAKE_CASE_,      keep_in_memory=SCREAMING_SNAKE_CASE_,      streaming=SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_,      )
          UpperCamelCase   :  Any						     =     load_from_cache_file
          UpperCamelCase   :  Any						     =     file_format
          UpperCamelCase   :  str						     =     Spark(
              df=SCREAMING_SNAKE_CASE_,      features=SCREAMING_SNAKE_CASE_,      cache_dir=SCREAMING_SNAKE_CASE_,      working_dir=SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_,      )






       def  snake_case_	(      self	)					->		Tuple:
          if self.streaming:
             return self.builder.as_streaming_dataset(split=self.split	)
          UpperCamelCase   :  Dict						     =     None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
          self.builder.download_and_prepare(
              download_mode=SCREAMING_SNAKE_CASE_,      file_format=self._file_format,      )
          return self.builder.as_dataset(split=self.split	)



