



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


_lowercase						  =		logging.get_logger(__name__)


_lowercase						  =		{"vocab_file": "vocab.json", "merges_file": "merges.txt"}

_lowercase						  =		{
    "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"
        ),
    },
}

_lowercase						  =		{
    "allenai/longformer-base-4096": 4096,
    "allenai/longformer-large-4096": 4096,
    "allenai/longformer-large-4096-finetuned-triviaqa": 4096,
    "allenai/longformer-base-4096-extra.pos.embd.only": 4096,
    "allenai/longformer-large-4096-extra.pos.embd.only": 4096,
}



@lru_cache()
# Copied from transformers.models.roberta.tokenization_roberta.bytes_to_unicode
def 		lowerCAmelCase__    (					)->  int:
    A__		 =					(
        list(range(ord('''!'''	) ,       ord('''~'''	) + 1	)	) + list(range(ord('''¡'''	) ,       ord('''¬'''	) + 1	)	) + list(range(ord('''®'''	) ,       ord('''ÿ'''	) + 1	)	)
    )
    A__		 =					bs[:]
    A__		 =					0
    for b in range(2**8	):
        if b not in bs:
            bs.append(a__	)
            cs.append(2**8 + n	)
            n += 1
    A__		 =					[chr(a__	) for n in cs]
    return dict(zip(a__ ,       a__	)	)



def 		lowerCAmelCase__    (					UpperCamelCase_				:       List[str]	)->  Tuple:
    A__		 =					set()
    A__		 =					word[0]
    for char in word[1:]:
        pairs.add((prev_char, char)	)
        A__		 =					char
    return pairs



class     _UpperCAmelCase			(       __lowercase     ):
     UpperCamelCase__		    =				VOCAB_FILES_NAMES
     UpperCamelCase__		    =				PRETRAINED_VOCAB_FILES_MAP
     UpperCamelCase__		    =				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     UpperCamelCase__		    =				['input_ids', 'attention_mask']




     def __init__(				self ,						a__ ,						a__ ,						a__="replace" ,						a__="<s>" ,						a__="</s>" ,						a__="</s>" ,						a__="<s>" ,						a__="<unk>" ,						a__="<pad>" ,						a__="<mask>" ,						a__=False ,						**a__ ,						):
         A__		 =					AddedToken(__A ,						lstrip=__A ,						rstrip=__A) if isinstance(__A ,						__A) else bos_token
         A__		 =					AddedToken(__A ,						lstrip=__A ,						rstrip=__A) if isinstance(__A ,						__A) else eos_token
         A__		 =					AddedToken(__A ,						lstrip=__A ,						rstrip=__A) if isinstance(__A ,						__A) else sep_token
         A__		 =					AddedToken(__A ,						lstrip=__A ,						rstrip=__A) if isinstance(__A ,						__A) else cls_token
         A__		 =					AddedToken(__A ,						lstrip=__A ,						rstrip=__A) if isinstance(__A ,						__A) else unk_token
         A__		 =					AddedToken(__A ,						lstrip=__A ,						rstrip=__A) if isinstance(__A ,						__A) else pad_token

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

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

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

         # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
         A__		 =					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):
         return len(self.encoder)




     def 		snake_case_		(				self):
         return dict(self.encoder ,						**self.added_tokens_encoder)




     def 		snake_case_		(				self ,						a__):
         if token in self.cache:
             return self.cache[token]
         A__		 =					tuple(__A)
         A__		 =					get_pairs(__A)

         if not pairs:
             return token

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

                 if word[i] == first and i < len(__A) - 1 and word[i + 1] == second:
                     new_word.append(first + second)
                     i += 2
                 else:
                     new_word.append(word[i])
                     i += 1
             A__		 =					tuple(__A)
             A__		 =					new_word
             if len(__A) == 1:
                 break
             else:
                 A__		 =					get_pairs(__A)
         A__		 =					''' '''.join(__A)
         A__		 =					word
         return word




     def 		snake_case_		(				self ,						a__):
         A__		 =					[]
         for token in re.findall(self.pat ,						__A):
             A__		 =					''''''.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(__A).split(''' '''))
         return bpe_tokens




     def 		snake_case_		(				self ,						a__):
         return self.encoder.get(__A ,						self.encoder.get(self.unk_token))




     def 		snake_case_		(				self ,						a__):
         return self.decoder.get(__A)




     def 		snake_case_		(				self ,						a__):
         A__		 =					''''''.join(__A)
         A__		 =					bytearray([self.byte_decoder[c] for c in text]).decode('''utf-8''' ,						errors=self.errors)
         return text




     def 		snake_case_		(				self ,						a__ ,						a__ = None):
         if not os.path.isdir(__A):
             logger.error(F"Vocabulary path ({save_directory}) should be a directory")
             return
         A__		 =					os.path.join(
             __A ,						(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file'''])
         A__		 =					os.path.join(
             __A ,						(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file'''])

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

         A__		 =					0
         with open(__A ,						'''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 a__: 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!''')
                     A__		 =					token_index
                 writer.write(''' '''.join(__A) + '''\n''')
                 index += 1

         return vocab_file, merge_file




     def 		snake_case_		(				self ,						a__ ,						a__ = None):
         if token_ids_a is None:
             return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
         A__		 =					[self.cls_token_id]
         A__		 =					[self.sep_token_id]
         return cls + token_ids_a + sep + sep + token_ids_a + sep




     def 		snake_case_		(				self ,						a__ ,						a__ = None ,						a__ = False):
         if already_has_special_tokens:
             return super().get_special_tokens_mask(
                 token_ids_a=__A ,						token_ids_a=__A ,						already_has_special_tokens=__A)

         if token_ids_a is None:
             return [1] + ([0] * len(__A)) + [1]
         return [1] + ([0] * len(__A)) + [1, 1] + ([0] * len(__A)) + [1]




     def 		snake_case_		(				self ,						a__ ,						a__ = None):
         A__		 =					[self.sep_token_id]
         A__		 =					[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 ,						a__ ,						a__=False ,						**a__):
         A__		 =					kwargs.pop('''add_prefix_space''' ,						self.add_prefix_space)
         if (is_split_into_words or add_prefix_space) and (len(__A) > 0 and not text[0].isspace()):
             A__		 =					''' ''' + text
         return (text, kwargs)





import jax.numpy as jnp

from ...utils import logging
from ..ta.modeling_flax_ta import FlaxTaEncoderModel, FlaxTaForConditionalGeneration, FlaxTaModel
from .configuration_mta import MTaConfig


_lowercase						  =		logging.get_logger(__name__)

_lowercase						  =		"T5Config"



def 		lowerCAmelCase__    (					UpperCamelCase_				:       jnp.array ,       UpperCamelCase_				:       int ,       UpperCamelCase_				:       int	)->  jnp.ndarray:
    A__		 =					jnp.zeros_like(UpperCamelCase_	)
    A__		 =					shifted_input_ids.at[:, 1:].set(input_ids[:, :-1]	)
    A__		 =					shifted_input_ids.at[:, 0].set(UpperCamelCase_	)

    A__		 =					jnp.where(shifted_input_ids == -1_0_0 ,       UpperCamelCase_ ,       UpperCamelCase_	)
    return shifted_input_ids



class     _UpperCAmelCase			(       A__     ):
     UpperCamelCase__		    =				'''mt5'''
     UpperCamelCase__		    =				MTaConfig







class     _UpperCAmelCase			(       A__     ):
     UpperCamelCase__		    =				'''mt5'''
     UpperCamelCase__		    =				MTaConfig







class     _UpperCAmelCase			(       A__     ):
     UpperCamelCase__		    =				'''mt5'''
     UpperCamelCase__		    =				MTaConfig
