




"""simple docstring"""


import json
import re
from typing import TYPE_CHECKING, List, Optional, Tuple, Union

import numpy as np

from ...utils import is_tf_available, is_torch_available, logging


if TYPE_CHECKING:
					if is_torch_available():
										import torch
					if is_tf_available():
										import tensorflow as tf

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_codegen import CodeGenTokenizer


A_      =    logging.get_logger(__name__)

A_      =    {'''vocab_file''': '''vocab.json''', '''merges_file''': '''merges.txt''', '''tokenizer_file''': '''tokenizer.json'''}

A_      =    {
    '''vocab_file''': {
        '''Salesforce/codegen-350M-mono''': '''https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/vocab.json''',
    },
    '''merges_file''': {
        '''Salesforce/codegen-350M-mono''': '''https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/merges.txt''',
    },
    '''tokenizer_file''': {
        '''Salesforce/codegen-350M-mono''': (
            '''https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/tokenizer.json'''
        ),
    },
}

A_      =    {
    '''Salesforce/codegen-350M-mono''': 20_48,
}






class  lowercase(					__a			):







	'''simple docstring'''





	lowercase__			 = VOCAB_FILES_NAMES
	lowercase__			 = PRETRAINED_VOCAB_FILES_MAP
	lowercase__			 = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	lowercase__			 = ["input_ids", "attention_mask"]
	lowercase__			 = CodeGenTokenizer


	def __init__(    self:		Union[str, Any],				a_:		Optional[int]=None,				a_:		Dict=None,				a_:		Any=None,				a_:		List[Any]="<|endoftext|>",				a_:		Dict="<|endoftext|>",				a_:		int="<|endoftext|>",				a_:		Dict=False,				**a_:		Optional[int],				):




			'''simple docstring'''



			super().__init__(
			    a_,				a_,				tokenizer_file=a_,				unk_token=a_,				bos_token=a_,				eos_token=a_,				add_prefix_space=a_,				**a_,				)

			if kwargs.pop("""add_bos_token""",				a_		):
					_snake_case							:		List[Any]								=  kwargs.pop("""name_or_path""",				""""""		)
					raise ValueError(
					    """Currenty GPT2's fast tokenizer does NOT support adding a BOS token."""
					    """Instead you should use GPT2's slow tokenizer class `CodeGenTokenizer` as follows: \n"""
					    f"`CodeGenTokenizer.from_pretrained('{model_id}')`\nor\n"
					    f"`AutoTokenizer.from_pretrained('{model_id}', use_fast=False)`\n"
					    """This issue will be fixed soon, see: https://github.com/huggingface/tokenizers/pull/1005."""
					    """ so that the fast tokenizer works correctly."""		)

			_snake_case							:		Dict								=  json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()		)
			if pre_tok_state.get("""add_prefix_space""",				a_		) != add_prefix_space:
					_snake_case							:		Optional[int]								=  getattr(a_,				pre_tok_state.pop("""type"""		)		)
					_snake_case							:		List[Any]								=  add_prefix_space
					_snake_case							:		Union[str, Any]								=  pre_tok_class(**a_		)

			_snake_case							:		Union[str, Any]								=  add_prefix_space


	def 			UpperCamelCase_       (    self:		Optional[Any],				*a_:		Any,				**a_:		Optional[int]		):




			'''simple docstring'''



			_snake_case							:		str								=  kwargs.get("""is_split_into_words""",				a_		)
			assert self.add_prefix_space or not is_split_into_words, (
			    f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
			    "to use it with pretokenized inputs."
			)

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


	def 			UpperCamelCase_       (    self:		Union[str, Any],				*a_:		Optional[int],				**a_:		Optional[int]		):




			'''simple docstring'''



			_snake_case							:		Union[str, Any]								=  kwargs.get("""is_split_into_words""",				a_		)

			assert self.add_prefix_space or not is_split_into_words, (
			    f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
			    "to use it with pretokenized inputs."
			)

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


	def 			UpperCamelCase_       (    self:		List[str],				a_:		str,				a_:		Optional[str] = None		):




			'''simple docstring'''



			_snake_case							:		List[Any]								=  self._tokenizer.model.save(a_,				name=a_		)
			return tuple(a_		)


	def 			UpperCamelCase_       (    self:		Optional[Any],				a_:		Union[int, List[int], "np.ndarray", "torch.Tensor", "tf.Tensor"],				a_:		bool = False,				a_:		bool = None,				a_:		Optional[List[str]] = None,				**a_:		Optional[int],				):




			'''simple docstring'''



			_snake_case							:		List[str]								=  super().decode(
			    token_ids=a_,				skip_special_tokens=a_,				clean_up_tokenization_spaces=a_,				**a_,				)

			if truncate_before_pattern is not None and len(a_		) > 0:
					_snake_case							:		str								=  self.truncate(a_,				a_		)

			return decoded_text


	def 			UpperCamelCase_       (    self:		Optional[Any],				a_:		int,				a_:		int		):




			'''simple docstring'''



			def find_re(a_:		str,				a_:		List[str],				a_:		Any		):
					_snake_case							:		List[Any]								=  pattern.search(a_,				a_		)
					return m.start() if m else -1

			_snake_case							:		List[Any]								=  [re.compile(a_,				re.MULTILINE		) for pattern in truncate_before_pattern]

			_snake_case							:		Union[str, Any]								=  list(re.finditer("""^print""",				a_,				re.MULTILINE		)		)

			if len(a_		) > 1:
					_snake_case							:		Dict								=  completion[: prints[1].start()]

			_snake_case							:		List[Any]								=  list(re.finditer("""^def""",				a_,				re.MULTILINE		)		)

			if len(a_		) > 1:
					_snake_case							:		Optional[int]								=  completion[: defs[1].start()]

			_snake_case							:		Dict								=  0

			_snake_case							:		str								=  [
			    pos for pos in [find_re(a_,				a_,				a_		) for terminal in terminals] if pos != -1
			]

			if len(a_		) > 0:
					return completion[: min(a_		)]
			else:
					return completion








"""simple docstring"""


import torch
from torch import nn






class  lowercase(					nn.Module			):







 '''simple docstring'''





 def __init__(    self:		Any,				a_:		List[str],				a_:		Union[str, Any],				a_:		int,				a_:		int,				a_:		List[Any]=1,				a_:		Union[str, Any]=False		):




   '''simple docstring'''



   super().__init__()

   _snake_case							:		int								=  n_token
   _snake_case							:		Tuple								=  d_embed
   _snake_case							:		List[str]								=  d_proj

   _snake_case							:		Optional[int]								=  cutoffs + [n_token]
   _snake_case							:		Any								=  [0] + self.cutoffs
   _snake_case							:		Tuple								=  div_val

   _snake_case							:		Optional[int]								=  self.cutoffs[0]
   _snake_case							:		Union[str, Any]								=  len(self.cutoffs		) - 1
   _snake_case							:		Union[str, Any]								=  self.shortlist_size + self.n_clusters

   if self.n_clusters > 0:
     _snake_case							:		List[Any]								=  nn.Parameter(torch.zeros(self.n_clusters,				self.d_embed		)		)
     _snake_case							:		Tuple								=  nn.Parameter(torch.zeros(self.n_clusters		)		)

   _snake_case							:		Any								=  nn.ModuleList()
   _snake_case							:		Optional[Any]								=  nn.ParameterList()

   if div_val == 1:
     for i in range(len(self.cutoffs		)		):
       if d_proj != d_embed:
         self.out_projs.append(nn.Parameter(torch.FloatTensor(a_,				a_		)		)		)
       else:
         self.out_projs.append(a_		)

     self.out_layers.append(nn.Linear(a_,				a_		)		)
   else:
     for i in range(len(self.cutoffs		)		):
       _snake_case		, _snake_case							:		List[Any]								=  self.cutoff_ends[i], self.cutoff_ends[i + 1]
       _snake_case							:		Union[str, Any]								=  d_embed // (div_val**i)

       self.out_projs.append(nn.Parameter(torch.FloatTensor(a_,				a_		)		)		)

       self.out_layers.append(nn.Linear(a_,				r_idx - l_idx		)		)

   _snake_case							:		Optional[int]								=  keep_order


 def 			UpperCamelCase_       (    self:		str,				a_:		Union[str, Any],				a_:		Dict,				a_:		int,				a_:		Tuple		):




   '''simple docstring'''



   if proj is None:
     _snake_case							:		List[Any]								=  nn.functional.linear(a_,				a_,				bias=a_		)
   else:
     # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
     _snake_case							:		List[Any]								=  nn.functional.linear(a_,				proj.t().contiguous()		)
     _snake_case							:		Tuple								=  nn.functional.linear(a_,				a_,				bias=a_		)
     # else:
     #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
     #     if bias is not None:
     #         logit = logit + bias

   return logit


 def 			UpperCamelCase_       (    self:		Dict,				a_:		Dict,				a_:		str=None,				a_:		Union[str, Any]=False		):




   '''simple docstring'''



   if labels is not None:
     # Shift so that tokens < n predict n
     _snake_case							:		int								=  hidden[..., :-1, :].contiguous()
     _snake_case							:		List[Any]								=  labels[..., 1:].contiguous()
     _snake_case							:		str								=  hidden.view(-1,				hidden.size(-1		)		)
     _snake_case							:		Dict								=  labels.view(-1		)
     if hidden.size(0		) != labels.size(0		):
       raise RuntimeError("""Input and labels should have the same size in the batch dimension."""		)
   else:
     _snake_case							:		int								=  hidden.view(-1,				hidden.size(-1		)		)

   if self.n_clusters == 0:
     _snake_case							:		Tuple								=  self._compute_logit(a_,				self.out_layers[0].weight,				self.out_layers[0].bias,				self.out_projs[0]		)
     if labels is not None:
       _snake_case							:		Dict								=  labels != -100
       _snake_case							:		str								=  torch.zeros_like(a_,				dtype=hidden.dtype,				device=hidden.device		)
       _snake_case							:		str								=  (
           -nn.functional.log_softmax(a_,				dim=-1		)[mask].gather(1,				labels[mask].unsqueeze(1		)		).squeeze(1		)
       )
     else:
       _snake_case							:		Optional[int]								=  nn.functional.log_softmax(a_,				dim=-1		)
   else:
     # construct weights and biases
     _snake_case		, _snake_case							:		Optional[int]								=  [], []
     for i in range(len(self.cutoffs		)		):
       if self.div_val == 1:
         _snake_case		, _snake_case							:		Optional[int]								=  self.cutoff_ends[i], self.cutoff_ends[i + 1]
         _snake_case							:		List[Any]								=  self.out_layers[0].weight[l_idx:r_idx]
         _snake_case							:		Tuple								=  self.out_layers[0].bias[l_idx:r_idx]
       else:
         _snake_case							:		Optional[int]								=  self.out_layers[i].weight
         _snake_case							:		int								=  self.out_layers[i].bias

       if i == 0:
         _snake_case							:		List[str]								=  torch.cat([weight_i, self.cluster_weight],				dim=0		)
         _snake_case							:		int								=  torch.cat([bias_i, self.cluster_bias],				dim=0		)

       weights.append(a_		)
       biases.append(a_		)

     _snake_case		, _snake_case		, _snake_case							:		Any								=  weights[0], biases[0], self.out_projs[0]

     _snake_case							:		List[str]								=  self._compute_logit(a_,				a_,				a_,				a_		)
     _snake_case							:		Union[str, Any]								=  nn.functional.log_softmax(a_,				dim=1		)

     if labels is None:
       _snake_case							:		Tuple								=  hidden.new_empty((head_logit.size(0		), self.n_token)		)
     else:
       _snake_case							:		Dict								=  torch.zeros_like(a_,				dtype=hidden.dtype,				device=hidden.device		)

     _snake_case							:		Union[str, Any]								=  0
     _snake_case							:		Optional[Any]								=  [0] + self.cutoffs
     for i in range(len(a_		) - 1		):
       _snake_case		, _snake_case							:		Dict								=  cutoff_values[i], cutoff_values[i + 1]

       if labels is not None:
         _snake_case							:		Dict								=  (labels >= l_idx) & (labels < r_idx)
         _snake_case							:		int								=  mask_i.nonzero().squeeze()

         if indices_i.numel() == 0:
           continue

         _snake_case							:		List[str]								=  labels.index_select(0,				a_		) - l_idx
         _snake_case							:		List[str]								=  head_logprob.index_select(0,				a_		)
         _snake_case							:		List[str]								=  hidden.index_select(0,				a_		)
       else:
         _snake_case							:		List[str]								=  hidden

       if i == 0:
         if labels is not None:
           _snake_case							:		Dict								=  head_logprob_i.gather(1,				target_i[:, None]		).squeeze(1		)
         else:
           _snake_case							:		Optional[Any]								=  head_logprob[:, : self.cutoffs[0]]
       else:
         _snake_case		, _snake_case		, _snake_case							:		Dict								=  weights[i], biases[i], self.out_projs[i]

         _snake_case							:		int								=  self._compute_logit(a_,				a_,				a_,				a_		)
         _snake_case							:		int								=  nn.functional.log_softmax(a_,				dim=1		)
         _snake_case							:		Dict								=  self.cutoffs[0] + i - 1  # No probability for the head cluster
         if labels is not None:
           _snake_case							:		Dict								=  head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
               1,				target_i[:, None]		).squeeze(1		)
         else:
           _snake_case							:		Optional[Any]								=  head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
           _snake_case							:		Any								=  logprob_i

       if labels is not None:
         if (hasattr(self,				"""keep_order"""		) and self.keep_order) or keep_order:
           out.index_copy_(0,				a_,				-logprob_i		)
         else:
           out[offset : offset + logprob_i.size(0		)].copy_(-logprob_i		)
         offset += logprob_i.size(0		)

   return out


 def 			UpperCamelCase_       (    self:		Union[str, Any],				a_:		Optional[int]		):




   '''simple docstring'''



   if self.n_clusters == 0:
     _snake_case							:		Optional[int]								=  self._compute_logit(a_,				self.out_layers[0].weight,				self.out_layers[0].bias,				self.out_projs[0]		)
     return nn.functional.log_softmax(a_,				dim=-1		)
   else:
     # construct weights and biases
     _snake_case		, _snake_case							:		List[Any]								=  [], []
     for i in range(len(self.cutoffs		)		):
       if self.div_val == 1:
         _snake_case		, _snake_case							:		Optional[Any]								=  self.cutoff_ends[i], self.cutoff_ends[i + 1]
         _snake_case							:		Optional[int]								=  self.out_layers[0].weight[l_idx:r_idx]
         _snake_case							:		str								=  self.out_layers[0].bias[l_idx:r_idx]
       else:
         _snake_case							:		List[Any]								=  self.out_layers[i].weight
         _snake_case							:		Union[str, Any]								=  self.out_layers[i].bias

       if i == 0:
         _snake_case							:		int								=  torch.cat([weight_i, self.cluster_weight],				dim=0		)
         _snake_case							:		Dict								=  torch.cat([bias_i, self.cluster_bias],				dim=0		)

       weights.append(a_		)
       biases.append(a_		)

     _snake_case		, _snake_case		, _snake_case							:		int								=  weights[0], biases[0], self.out_projs[0]
     _snake_case							:		List[Any]								=  self._compute_logit(a_,				a_,				a_,				a_		)

     _snake_case							:		List[Any]								=  hidden.new_empty((head_logit.size(0		), self.n_token)		)
     _snake_case							:		int								=  nn.functional.log_softmax(a_,				dim=1		)

     _snake_case							:		List[Any]								=  [0] + self.cutoffs
     for i in range(len(a_		) - 1		):
       _snake_case		, _snake_case							:		List[str]								=  cutoff_values[i], cutoff_values[i + 1]

       if i == 0:
         _snake_case							:		List[str]								=  head_logprob[:, : self.cutoffs[0]]
       else:
         _snake_case		, _snake_case		, _snake_case							:		List[Any]								=  weights[i], biases[i], self.out_projs[i]

         _snake_case							:		Optional[int]								=  self._compute_logit(a_,				a_,				a_,				a_		)
         _snake_case							:		Any								=  nn.functional.log_softmax(a_,				dim=1		)

         _snake_case							:		Dict								=  head_logprob[:, -i] + tail_logprob_i
         _snake_case							:		Any								=  logprob_i

     return out


