



import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


UpperCAmelCase__       =	get_logger(__name__)


UpperCAmelCase__       =	r"\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n"







class       __lowerCAmelCase   :



			@add_start_docstrings(A)
			def __call__(						self       :  Optional[int]			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray)       ->   jnp.ndarray:



						"""simple docstring"""






						raise NotImplementedError(
						    F"{self.__class__} is an abstract class. Only classes inheriting this class can be called.")







class       __lowerCAmelCase   :



			@add_start_docstrings(A)
			def __call__(						self       :  List[Any]			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray)       ->   jnp.ndarray:



						"""simple docstring"""






						raise NotImplementedError(
						    F"{self.__class__} is an abstract class. Only classes inheriting this class can be called.")







class       __lowerCAmelCase   (      A						):



			@add_start_docstrings(A)
			def __call__(						self       :  Union[str, Any]			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray			,		A       :  int			,		**A       :  List[Any])       ->   jnp.ndarray:



						"""simple docstring"""






						for processor in self:
									_UpperCAmelCase			=       inspect.signature(processor.__call__).parameters
									if len(A) > 3:
												if not all(arg in kwargs for arg in list(function_args.keys())[2:]):
															raise ValueError(
															    F"Make sure that all the required parameters: {list(function_args.keys())} for "
															    F"{processor.__class__} are passed to the logits processor.")
												_UpperCAmelCase			=       processor(A			,		A			,		A			,		**A)
									else:
												_UpperCAmelCase			=       processor(A			,		A			,		A)
						return scores







class       __lowerCAmelCase   (      A						):



			def __init__(						self       :  List[str]			,		A       :  float)       ->   Optional[int]:



						"""simple docstring"""






						if not isinstance(A			,		A) or not (temperature > 0):
									raise ValueError(F"`temperature` has to be a strictly positive float, but is {temperature}")

						_UpperCAmelCase			=       temperature



			def __call__(						self       :  Union[str, Any]			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray			,		A       :  int)       ->   jnp.ndarray:



						"""simple docstring"""






						_UpperCAmelCase			=       scores / self.temperature
						return scores







class       __lowerCAmelCase   (      A						):



			def __init__(						self       :  Union[str, Any]			,		A       :  float			,		A       :  float = -float('Inf')			,		A       :  int = 1)       ->   str:



						"""simple docstring"""






						if not isinstance(A			,		A) or (top_p < 0 or top_p > 1.0):
									raise ValueError(F"`top_p` has to be a float > 0 and < 1, but is {top_p}")
						if not isinstance(A			,		A) or (min_tokens_to_keep < 1):
									raise ValueError(F"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}")

						_UpperCAmelCase			=       top_p
						_UpperCAmelCase			=       filter_value
						_UpperCAmelCase			=       min_tokens_to_keep



			def __call__(						self       :  Optional[Any]			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray			,		A       :  int)       ->   jnp.ndarray:



						"""simple docstring"""






						_UpperCAmelCase	,					_UpperCAmelCase			=       lax.top_k(A			,		scores.shape[-1])

						_UpperCAmelCase			=       jnp.full_like(A			,		self.filter_value)
						_UpperCAmelCase			=       jax.nn.softmax(A			,		axis=-1).cumsum(axis=-1)
						_UpperCAmelCase			=       cumulative_probs < self.top_p

						# include the token that is higher than top_p as well
						_UpperCAmelCase			=       jnp.roll(A			,		1)
						score_mask |= score_mask.at[:, 0].set(A)

						# min tokens to keep
						_UpperCAmelCase			=       score_mask.at[:, : self.min_tokens_to_keep].set(A)

						_UpperCAmelCase			=       jnp.where(A			,		A			,		A)
						_UpperCAmelCase			=       jax.lax.sort_key_val(A			,		A)[-1]

						return next_scores







class       __lowerCAmelCase   (      A						):



			def __init__(						self       :  Any			,		A       :  int			,		A       :  float = -float('Inf')			,		A       :  int = 1)       ->   Optional[Any]:



						"""simple docstring"""






						if not isinstance(A			,		A) or top_k <= 0:
									raise ValueError(F"`top_k` has to be a strictly positive integer, but is {top_k}")

						_UpperCAmelCase			=       max(A			,		A)
						_UpperCAmelCase			=       filter_value



			def __call__(						self       :  Any			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray			,		A       :  int)       ->   jnp.ndarray:



						"""simple docstring"""






						_UpperCAmelCase	,					_UpperCAmelCase			=       scores.shape
						_UpperCAmelCase			=       jnp.full(batch_size * vocab_size			,		self.filter_value)

						_UpperCAmelCase			=       min(self.top_k			,		scores.shape[-1])  # Safety check
						_UpperCAmelCase	,					_UpperCAmelCase			=       lax.top_k(A			,		A)
						_UpperCAmelCase			=       jnp.broadcast_to((jnp.arange(A) * vocab_size)[:, None]			,		(batch_size, topk)).flatten()
						_UpperCAmelCase			=       topk_scores.flatten()
						_UpperCAmelCase			=       topk_indices.flatten() + shift

						_UpperCAmelCase			=       next_scores_flat.at[topk_indices_flat].set(A)
						_UpperCAmelCase			=       next_scores_flat.reshape(A			,		A)
						return next_scores







class       __lowerCAmelCase   (      A						):



			def __init__(						self       :  List[str]			,		A       :  int)       ->   Any:



						"""simple docstring"""






						_UpperCAmelCase			=       bos_token_id



			def __call__(						self       :  List[str]			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray			,		A       :  int)       ->   jnp.ndarray:



						"""simple docstring"""






						_UpperCAmelCase			=       jnp.full(scores.shape			,		-float('inf'))

						_UpperCAmelCase			=       1 - jnp.bool_(cur_len - 1)

						_UpperCAmelCase			=       jnp.where(A			,		new_scores.at[:, self.bos_token_id].set(0)			,		A)

						return scores







class       __lowerCAmelCase   (      A						):



			def __init__(						self       :  List[Any]			,		A       :  int			,		A       :  int)       ->   Dict:



						"""simple docstring"""






						_UpperCAmelCase			=       max_length
						_UpperCAmelCase			=       eos_token_id



			def __call__(						self       :  int			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray			,		A       :  int)       ->   jnp.ndarray:



						"""simple docstring"""






						_UpperCAmelCase			=       jnp.full(scores.shape			,		-float('inf'))

						_UpperCAmelCase			=       1 - jnp.bool_(cur_len - self.max_length + 1)

						_UpperCAmelCase			=       jnp.where(A			,		new_scores.at[:, self.eos_token_id].set(0)			,		A)

						return scores







class       __lowerCAmelCase   (      A						):



			def __init__(						self       :  Union[str, Any]			,		A       :  int			,		A       :  int)       ->   Tuple:



						"""simple docstring"""






						if not isinstance(A			,		A) or min_length < 0:
									raise ValueError(F"`min_length` has to be a positive integer, but is {min_length}")

						if not isinstance(A			,		A) or eos_token_id < 0:
									raise ValueError(F"`eos_token_id` has to be a positive integer, but is {eos_token_id}")

						_UpperCAmelCase			=       min_length
						_UpperCAmelCase			=       eos_token_id



			def __call__(						self       :  List[str]			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray			,		A       :  int)       ->   jnp.ndarray:



						"""simple docstring"""






						_UpperCAmelCase			=       1 - jnp.clip(cur_len - self.min_length			,		0			,		1)

						_UpperCAmelCase			=       jnp.where(A			,		scores.at[:, self.eos_token_id].set(-float('inf'))			,		A)

						return scores







class       __lowerCAmelCase   (      A						):



			def __init__(						self       :  Dict			,		A       :  List[Any]			,		A       :  List[Any])       ->   Any:



						"""simple docstring"""






						_UpperCAmelCase			=       list(A)
						_UpperCAmelCase			=       begin_index



			def __call__(						self       :  Any			,		A       :  Optional[Any]			,		A       :  Tuple			,		A       :  int)       ->   Optional[Any]:



						"""simple docstring"""






						_UpperCAmelCase			=       1 - jnp.bool_(cur_len - self.begin_index)

						_UpperCAmelCase			=       jnp.where(A			,		scores.at[:, self.begin_suppress_tokens].set(-float('inf'))			,		A)

						return scores







class       __lowerCAmelCase   (      A						):



			def __init__(						self       :  Dict			,		A       :  list)       ->   List[Any]:



						"""simple docstring"""






						_UpperCAmelCase			=       list(A)



			def __call__(						self       :  Union[str, Any]			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray			,		A       :  int)       ->   jnp.ndarray:



						"""simple docstring"""






						_UpperCAmelCase			=       scores.at[..., self.suppress_tokens].set(-float('inf'))

						return scores







class       __lowerCAmelCase   (      A						):



			def __init__(						self       :  Optional[int]			,		A       :  Dict)       ->   int:



						"""simple docstring"""






						_UpperCAmelCase			=       dict(A)
						# Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
						# index of the array corresponds to the index of the token to be forced, for XLA compatibility.
						# Indexes without forced tokens will have a negative value.
						_UpperCAmelCase			=       jnp.ones((max(force_token_map.keys()) + 1)			,		dtype=jnp.intaa) * -1
						for index, token in force_token_map.items():
									if token is not None:
												_UpperCAmelCase			=       force_token_array.at[index].set(A)
						_UpperCAmelCase			=       jnp.intaa(A)



			def __call__(						self       :  Union[str, Any]			,		A       :  jnp.ndarray			,		A       :  jnp.ndarray			,		A       :  int)       ->   jnp.ndarray:



						"""simple docstring"""






						def _force_token(A       :  Dict):
									_UpperCAmelCase			=       scores.shape[0]
									_UpperCAmelCase			=       self.force_token_array[generation_idx]

									_UpperCAmelCase			=       jnp.ones_like(A			,		dtype=scores.dtype) * -float('inf')
									_UpperCAmelCase			=       jnp.zeros((batch_size, 1)			,		dtype=scores.dtype)
									_UpperCAmelCase			=       lax.dynamic_update_slice(A			,		A			,		(0, current_token))
									return new_scores

						_UpperCAmelCase			=       lax.cond(
						    cur_len >= self.force_token_array.shape[0]			,		lambda: scores			,		lambda: lax.cond(
						        self.force_token_array[cur_len] >= 0			,		lambda: _force_token(A)			,		lambda: scores			,		)			,		)
						return scores







class       __lowerCAmelCase   (      A						):



			def __init__(						self       :  int			,		A       :  str			,		A       :  List[Any]			,		A       :  List[Any])       ->   List[Any]:



						"""simple docstring"""






						_UpperCAmelCase			=       generate_config.eos_token_id
						_UpperCAmelCase			=       generate_config.no_timestamps_token_id
						_UpperCAmelCase			=       generate_config.no_timestamps_token_id + 1

						_UpperCAmelCase			=       decoder_input_length + 1

						if generate_config.is_multilingual:
									# room for language token and task token
									self.begin_index += 2
						if hasattr(A			,		'max_initial_timestamp_index'):
									_UpperCAmelCase			=       generate_config.max_initial_timestamp_index
						else:
									_UpperCAmelCase			=       model_config.vocab_size
						if self.max_initial_timestamp_index is None:
									_UpperCAmelCase			=       model_config.vocab_size



			def __call__(						self       :  List[str]			,		A       :  Tuple			,		A       :  Optional[int]			,		A       :  Optional[int])       ->   Any:



						"""simple docstring"""






						_UpperCAmelCase			=       scores.at[:, self.no_timestamps_token_id].set(-float('inf'))

						def handle_pairs(A       :  Optional[Any]			,		A       :  Optional[Any]):
									_UpperCAmelCase			=       jnp.where((cur_len - self.begin_index) >= 1			,		A			,		A)
									_UpperCAmelCase			=       jnp.where(
									    input_ids_k[cur_len - 1] >= self.timestamp_begin			,		True and last_was_timestamp			,		A			,		)

									_UpperCAmelCase			=       jnp.where((cur_len - self.begin_index) < 2			,		A			,		A)
									_UpperCAmelCase			=       jnp.where(
									    input_ids_k[cur_len - 2] >= self.timestamp_begin			,		A			,		A			,		)

									return jnp.where(
									    A			,		jnp.where(
									        penultimate_was_timestamp > 0			,		scores_k.at[self.timestamp_begin :].set(-float('inf'))			,		scores_k.at[: self.eos_token_id].set(-float('inf'))			,		)			,		A			,		)

						_UpperCAmelCase			=       jax.vmap(A)(A			,		A)

						_UpperCAmelCase			=       jnp.where(cur_len == self.begin_index			,		A			,		A)
						_UpperCAmelCase			=       jnp.where(
						    self.max_initial_timestamp_index is not None			,		True and apply_max_initial_timestamp			,		A			,		)

						_UpperCAmelCase			=       self.timestamp_begin + self.max_initial_timestamp_index

						_UpperCAmelCase			=       jnp.where(
						    A			,		scores.at[:, last_allowed + 1 :].set(-float('inf'))			,		A			,		)

						# if sum of probability over timestamps is above any other token, sample timestamp
						_UpperCAmelCase			=       jax.nn.log_softmax(A			,		axis=-1)

						def handle_cumulative_probs(A       :  Union[str, Any]			,		A       :  str):
									_UpperCAmelCase			=       jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]			,		axis=-1)
									_UpperCAmelCase			=       jnp.max(logprobs_k[: self.timestamp_begin])
									return jnp.where(
									    timestamp_logprob > max_text_token_logprob			,		scores_k.at[: self.timestamp_begin].set(-float('inf'))			,		A			,		)

						_UpperCAmelCase			=       jax.vmap(A)(A			,		A)

						return scores





import json
import os
import unittest

from transformers import CLIPTokenizer, CLIPTokenizerFast
from transformers.models.clip.tokenization_clip import VOCAB_FILES_NAMES
from transformers.testing_utils import require_ftfy, require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin







@require_tokenizers
class       __lowerCAmelCase   (      A    ,						unittest.TestCase						):
			UpperCamelCase        =					CLIPTokenizer
			UpperCamelCase        =					CLIPTokenizerFast
			UpperCamelCase        =					True
			UpperCamelCase        =					{}
			UpperCamelCase        =					False



			def    _lowerCamelCase	(						self       :  List[str])       ->   List[str]:



						"""simple docstring"""






						super().setUp()

						# fmt: off
						_UpperCAmelCase			=       ['l', 'o', 'w', 'e', 'r', 's', 't', 'i', 'd', 'n', 'lo', 'l</w>', 'w</w>', 'r</w>', 't</w>', 'low</w>', 'er</w>', 'lowest</w>', 'newer</w>', 'wider', '<unk>', '<|startoftext|>', '<|endoftext|>']
						# fmt: on
						_UpperCAmelCase			=       dict(zip(A			,		range(len(A))))
						_UpperCAmelCase			=       ['#version: 0.2', 'l o', 'lo w</w>', 'e r</w>']
						_UpperCAmelCase			=       {'unk_token': '<unk>'}

						_UpperCAmelCase			=       os.path.join(self.tmpdirname			,		VOCAB_FILES_NAMES['vocab_file'])
						_UpperCAmelCase			=       os.path.join(self.tmpdirname			,		VOCAB_FILES_NAMES['merges_file'])
						with open(self.vocab_file			,		'w'			,		encoding='utf-8') as fp:
									fp.write(json.dumps(A) + '\n')
						with open(self.merges_file			,		'w'			,		encoding='utf-8') as fp:
									fp.write('\n'.join(A))



			def    _lowerCamelCase	(						self       :  Optional[Any]			,		**A       :  str)       ->   Optional[int]:



						"""simple docstring"""






						kwargs.update(self.special_tokens_map)
						return CLIPTokenizer.from_pretrained(self.tmpdirname			,		**A)



			def    _lowerCamelCase	(						self       :  Any			,		**A       :  Dict)       ->   Union[str, Any]:



						"""simple docstring"""






						kwargs.update(self.special_tokens_map)
						return CLIPTokenizerFast.from_pretrained(self.tmpdirname			,		**A)



			def    _lowerCamelCase	(						self       :  Optional[int]			,		A       :  Union[str, Any])       ->   int:



						"""simple docstring"""






						_UpperCAmelCase			=       'lower newer'
						_UpperCAmelCase			=       'lower newer'
						return input_text, output_text



			def    _lowerCamelCase	(						self       :  Dict)       ->   Any:



						"""simple docstring"""






						_UpperCAmelCase			=       CLIPTokenizer(self.vocab_file			,		self.merges_file			,		**self.special_tokens_map)
						_UpperCAmelCase			=       'lower newer'
						_UpperCAmelCase			=       ['lo', 'w', 'er</w>', 'n', 'e', 'w', 'er</w>']
						_UpperCAmelCase			=       tokenizer.tokenize(A)
						self.assertListEqual(A			,		A)

						_UpperCAmelCase			=       tokens + [tokenizer.unk_token]
						_UpperCAmelCase			=       [10, 2, 16, 9, 3, 2, 16, 20]
						self.assertListEqual(tokenizer.convert_tokens_to_ids(A)			,		A)



			@require_ftfy
			def    _lowerCamelCase	(						self       :  List[Any])       ->   Dict:



						"""simple docstring"""






						for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
									with self.subTest(F"{tokenizer.__class__.__name__} ({pretrained_name})"):
												_UpperCAmelCase			=       self.tokenizer_class.from_pretrained(A			,		**A)
												_UpperCAmelCase			=       self.rust_tokenizer_class.from_pretrained(A			,		**A)

												_UpperCAmelCase			=       'A\n\'ll 11p223RF☆ho!!to?\'d\'d\'\'d of a cat to-$\'\'d.'
												_UpperCAmelCase			=       tokenizer_s.tokenize(A)
												_UpperCAmelCase			=       tokenizer_r.tokenize(A)

												self.assertListEqual(A			,		A)

												# Test that the tokenization is identical on an example containing a character (Latin Small Letter A
												# with Tilde) encoded in 2 different ways
												_UpperCAmelCase			=       'xa\u0303y' + ' ' + 'x\xe3y'
												_UpperCAmelCase			=       tokenizer_s.tokenize(A)
												_UpperCAmelCase			=       tokenizer_r.tokenize(A)

												self.assertListEqual(A			,		A)

												# Test that the tokenization is identical on unicode of space type
												_UpperCAmelCase			=       [
												    '\u0009',  # (horizontal tab, '\t')
												    '\u000B',  # (vertical tab)
												    '\u000C',  # (form feed)
												    '\u0020',  # (space, ' ')
												    '\u200E',  # (left-to-right mark):w
												    '\u200F',  # (right-to-left mark)
												]
												for unicode_seq in spaces_unicodes:
															_UpperCAmelCase			=       tokenizer_s.tokenize(A)
															_UpperCAmelCase			=       tokenizer_r.tokenize(A)

															self.assertListEqual(A			,		A)

												# Test that the tokenization is identical on unicode of line break type
												_UpperCAmelCase			=       [
												    '\u000A',  # (line feed, '\n')
												    '\r\n',  # (carriage return and line feed, '\r\n')
												    '\u000D',  # (carriage return, '\r')
												    '\r',  # (carriage return, '\r')
												    '\u000D',  # (carriage return, '\r')
												    '\u2028',  # (line separator)
												    '\u2029',  # (paragraph separator)
												    # "\u0085", # (next line)
												]

												# The tokenization is not identical for the character "\u0085" (next line). The slow version using ftfy transforms
												# it into the Horizontal Ellipsis character "…" ("\u2026") while the fast version transforms it into a
												# space (and thus into an empty list).

												for unicode_seq in line_break_unicodes:
															_UpperCAmelCase			=       tokenizer_s.tokenize(A)
															_UpperCAmelCase			=       tokenizer_r.tokenize(A)

															self.assertListEqual(A			,		A)



			def    _lowerCamelCase	(						self       :  str)       ->   Optional[Any]:



						"""simple docstring"""






						for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
									with self.subTest(F"{tokenizer.__class__.__name__} ({pretrained_name})"):
												_UpperCAmelCase			=       'hello'  # `hello` is a token in the vocabulary of `pretrained_name`
												_UpperCAmelCase			=       F"{text_of_1_token} {text_of_1_token}"

												_UpperCAmelCase			=       self.rust_tokenizer_class.from_pretrained(
												    A			,		use_fast=A			,		)
												_UpperCAmelCase			=       tokenizer_r(A			,		return_offsets_mapping=A			,		add_special_tokens=A)
												self.assertEqual(encoding.offset_mapping[0]			,		(0, len(A)))
												self.assertEqual(
												    encoding.offset_mapping[1]			,		(len(A) + 1, len(A) + 1 + len(A))			,		)

												_UpperCAmelCase			=       F" {text}"

												_UpperCAmelCase			=       self.rust_tokenizer_class.from_pretrained(
												    A			,		use_fast=A			,		)
												_UpperCAmelCase			=       tokenizer_r(A			,		return_offsets_mapping=A			,		add_special_tokens=A)
												self.assertEqual(encoding.offset_mapping[0]			,		(1, 1 + len(A)))
												self.assertEqual(
												    encoding.offset_mapping[1]			,		(1 + len(A) + 1, 1 + len(A) + 1 + len(A))			,		)



			def    _lowerCamelCase	(						self       :  Tuple)       ->   str:



						"""simple docstring"""






						with self.assertRaises(A) as context:
									self.rust_tokenizer_class.from_pretrained('robot-test/old-clip-tokenizer')

						self.assertTrue(
						    context.exception.args[0].startswith(
						        'The `backend_tokenizer` provided does not match the expected format.'))



			@require_ftfy
			def    _lowerCamelCase	(						self       :  int)       ->   int:



						"""simple docstring"""






						super().test_tokenization_python_rust_equals()



			def    _lowerCamelCase	(						self       :  Union[str, Any])       ->   Any:



						"""simple docstring"""






						pass
