
from typing import Dict

import numpy as np

from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, GenericTensor, Pipeline, PipelineException


if is_tf_available():
			import tensorflow as tf

			from ..tf_utils import stable_softmax


if is_torch_available():
			import torch


_UpperCAmelCase      :				Union[str, Any]          = logging.get_logger(__name__)


@add_end_docstrings(
    __lowerCAmelCase						,     r"\n        top_k (`int`, defaults to 5):\n            The number of predictions to return.\n        targets (`str` or `List[str]`, *optional*):\n            When passed, the model will limit the scores to the passed targets instead of looking up in the whole\n            vocab. If the provided targets are not in the model vocab, they will be tokenized and the first resulting\n            token will be used (with a warning, and that might be slower).\n\n    "						,     )
class        lowercase      (	__lowerCAmelCase     ):



		def __UpperCamelCase				(				self       ,   A_ )						->      np.ndarray:
				"""simple docstring"""
				if self.framework == "tf":
						UpperCamelCase						=  tf.where(input_ids == self.tokenizer.mask_token_id ).numpy()
				elif self.framework == "pt":
						UpperCamelCase						=  torch.nonzero(input_ids == self.tokenizer.mask_token_id       ,   as_tuple=lowerCAmelCase_ )
				else:
						raise ValueError('Unsupported framework' )
				return masked_index



		def __UpperCamelCase				(				self       ,   A_ )						->      np.ndarray:
				"""simple docstring"""
				UpperCamelCase						=  self.get_masked_index(lowerCAmelCase_ )
				UpperCamelCase						=  np.prod(masked_index.shape )
				if numel < 1:
						raise PipelineException(
						    'fill-mask'       ,   self.model.base_model_prefix       ,   F'''No mask_token ({self.tokenizer.mask_token}) found on the input'''       ,   )



		def __UpperCamelCase				(				self       ,   A_ )						->      int:
				"""simple docstring"""
				if isinstance(lowerCAmelCase_       ,   lowerCAmelCase_ ):
						for model_input in model_inputs:
								self._ensure_exactly_one_mask_token(model_input['input_ids'][0] )
				else:
						for input_ids in model_inputs["input_ids"]:
								self._ensure_exactly_one_mask_token(lowerCAmelCase_ )



		def __UpperCamelCase				(				self       ,   A_       ,   A_=None       ,   **A_ )						->      Dict[str, GenericTensor]:
				"""simple docstring"""
				if return_tensors is None:
						UpperCamelCase						=  self.framework
				UpperCamelCase						=  self.tokenizer(lowerCAmelCase_       ,   return_tensors=lowerCAmelCase_ )
				self.ensure_exactly_one_mask_token(lowerCAmelCase_ )
				return model_inputs



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.model(**lowerCAmelCase_ )
				UpperCamelCase						=  model_inputs['input_ids']
				return model_outputs



		def __UpperCamelCase				(				self       ,   A_       ,   A_=5       ,   A_=None )						->      Optional[int]:
				"""simple docstring"""
				# Cap top_k if there are targets
				if target_ids is not None and target_ids.shape[0] < top_k:
						UpperCamelCase						=  target_ids.shape[0]
				UpperCamelCase						=  model_outputs['input_ids'][0]
				UpperCamelCase						=  model_outputs['logits']

				if self.framework == "tf":
						UpperCamelCase						=  tf.where(input_ids == self.tokenizer.mask_token_id ).numpy()[:, 0]

						UpperCamelCase						=  outputs.numpy()

						UpperCamelCase						=  outputs[0, masked_index, :]
						UpperCamelCase						=  stable_softmax(lowerCAmelCase_       ,   axis=-1 )
						if target_ids is not None:
								UpperCamelCase						=  tf.gather_nd(tf.squeeze(lowerCAmelCase_       ,   0 )       ,   target_ids.reshape(-1       ,   1 ) )
								UpperCamelCase						=  tf.expand_dims(lowerCAmelCase_       ,   0 )

						UpperCamelCase						=  tf.math.top_k(lowerCAmelCase_       ,   k=lowerCAmelCase_ )
						UpperCamelCase			,						UpperCamelCase						=  topk.values.numpy(), topk.indices.numpy()
				else:
						UpperCamelCase						=  torch.nonzero(input_ids == self.tokenizer.mask_token_id       ,   as_tuple=lowerCAmelCase_ ).squeeze(-1 )
						# Fill mask pipeline supports only one ${mask_token} per sample

						UpperCamelCase						=  outputs[0, masked_index, :]
						UpperCamelCase						=  logits.softmax(dim=-1 )
						if target_ids is not None:
								UpperCamelCase						=  probs[..., target_ids]

						UpperCamelCase			,						UpperCamelCase						=  probs.topk(lowerCAmelCase_ )

				UpperCamelCase						=  []
				UpperCamelCase						=  values.shape[0] == 1
				for i, (_values, _predictions) in enumerate(zip(values.tolist()       ,   predictions.tolist() ) ):
						UpperCamelCase						=  []
						for v, p in zip(_values       ,   _predictions ):
								# Copy is important since we're going to modify this array in place
								UpperCamelCase						=  input_ids.numpy().copy()
								if target_ids is not None:
										UpperCamelCase						=  target_ids[p].tolist()

								UpperCamelCase						=  p
								# Filter padding out:
								UpperCamelCase						=  tokens[np.where(tokens != self.tokenizer.pad_token_id )]
								# Originally we skip special tokens to give readable output.
								# For multi masks though, the other [MASK] would be removed otherwise
								# making the output look odd, so we add them back
								UpperCamelCase						=  self.tokenizer.decode(lowerCAmelCase_       ,   skip_special_tokens=lowerCAmelCase_ )
								UpperCamelCase						=  {'score': v, 'token': p, 'token_str': self.tokenizer.decode([p] ), 'sequence': sequence}
								row.append(lowerCAmelCase_ )
						result.append(lowerCAmelCase_ )
				if single_mask:
						return result[0]
				return result



		def __UpperCamelCase				(				self       ,   A_       ,   A_=None )						->      Optional[Any]:
				"""simple docstring"""
				if isinstance(lowerCAmelCase_       ,   lowerCAmelCase_ ):
						UpperCamelCase						=  [targets]
				try:
						UpperCamelCase						=  self.tokenizer.get_vocab()
				except Exception:
						UpperCamelCase						=  {}
				UpperCamelCase						=  []
				for target in targets:
						UpperCamelCase						=  vocab.get(lowerCAmelCase_       ,   lowerCAmelCase_ )
						if id_ is None:
								UpperCamelCase						=  self.tokenizer(
								    lowerCAmelCase_       ,   add_special_tokens=lowerCAmelCase_       ,   return_attention_mask=lowerCAmelCase_       ,   return_token_type_ids=lowerCAmelCase_       ,   max_length=1       ,   truncation=lowerCAmelCase_       ,   )['input_ids']
								if len(lowerCAmelCase_ ) == 0:
										logger.warning(
										    F'''The specified target token `{target}` does not exist in the model vocabulary. '''
										    'We cannot replace it with anything meaningful, ignoring it' )
										continue
								UpperCamelCase						=  input_ids[0]
								# XXX: If users encounter this pass
								# it becomes pretty slow, so let's make sure
								# The warning enables them to fix the input to
								# get faster performance.
								logger.warning(
								    F'''The specified target token `{target}` does not exist in the model vocabulary. '''
								    F'''Replacing with `{self.tokenizer.convert_ids_to_tokens(id_ )}`.''' )
						target_ids.append(id_ )
				UpperCamelCase						=  list(set(lowerCAmelCase_ ) )
				if len(lowerCAmelCase_ ) == 0:
						raise ValueError('At least one target must be provided when passed.' )
				UpperCamelCase						=  np.array(lowerCAmelCase_ )
				return target_ids



		def __UpperCamelCase				(				self       ,   A_=None       ,   A_=None )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  {}

				if targets is not None:
						UpperCamelCase						=  self.get_target_ids(lowerCAmelCase_       ,   lowerCAmelCase_ )
						UpperCamelCase						=  target_ids

				if top_k is not None:
						UpperCamelCase						=  top_k

				if self.tokenizer.mask_token_id is None:
						raise PipelineException(
						    'fill-mask'       ,   self.model.base_model_prefix       ,   'The tokenizer does not define a `mask_token`.' )
				return {}, {}, postprocess_params



		def __call__(				self       ,   A_       ,   *A_       ,   **A_ )						->      Dict:
				"""simple docstring"""
				UpperCamelCase						=  super().__call__(lowerCAmelCase_       ,   **lowerCAmelCase_ )
				if isinstance(lowerCAmelCase_       ,   lowerCAmelCase_ ) and len(lowerCAmelCase_ ) == 1:
						return outputs[0]
				return outputs


from abc import ABC, abstractmethod
from typing import List, Optional


class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self )						->      Optional[Any]:
				"""simple docstring"""
				# test for the above condition
				self.test()



		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				UpperCamelCase						=  0
				UpperCamelCase						=  False
				while not completed:
						if counter == 1:
								self.reset()
						UpperCamelCase						=  self.advance()
						if not self.does_advance(A_ ):
								raise Exception(
								    'Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.' )

						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase						=  self.update(A_ )
						counter += 1

						if counter > 10_000:
								raise Exception('update() does not fulfill the constraint.' )

				if self.remaining() != 0:
						raise Exception('Custom Constraint is not defined correctly.' )



		@abstractmethod
		def __UpperCamelCase				(				self )						->      Optional[Any]:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self       ,   A_ )						->      str:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self       ,   A_ )						->      int:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self )						->      Any:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self )						->      str:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self       ,   A_=False )						->      int:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_ )						->      Any:
				"""simple docstring"""
				super(A_       ,   self ).__init__()

				if not isinstance(A_       ,   A_ ) or len(A_ ) == 0:
						raise ValueError(F'''`token_ids` has to be a non-empty list, but is {token_ids}.''' )
				if any((not isinstance(A_       ,   A_ ) or token_id < 0) for token_id in token_ids ):
						raise ValueError(F'''Each list in `token_ids` has to be a list of positive integers, but is {token_ids}.''' )

				UpperCamelCase						=  token_ids

				UpperCamelCase						=  len(self.token_ids )
				UpperCamelCase						=  -1  # the index of the currently fulfilled step
				UpperCamelCase						=  False



		def __UpperCamelCase				(				self )						->      Optional[Any]:
				"""simple docstring"""
				if self.completed:
						return None
				return self.token_ids[self.fulfilled_idx + 1]



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(A_ )}''' )

				if self.completed:
						return False

				return token_id == self.token_ids[self.fulfilled_idx + 1]



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(A_ )}''' )

				UpperCamelCase						=  False
				UpperCamelCase						=  False
				UpperCamelCase						=  False

				if self.does_advance(A_ ):
						self.fulfilled_idx += 1
						UpperCamelCase						=  True
						if self.fulfilled_idx == (self.seqlen - 1):
								UpperCamelCase						=  True
						UpperCamelCase						=  completed
				else:
						# failed to make progress.
						UpperCamelCase						=  True
						self.reset()
				return stepped, completed, reset



		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				UpperCamelCase						=  False
				UpperCamelCase						=  0



		def __UpperCamelCase				(				self )						->      int:
				"""simple docstring"""
				return self.seqlen - (self.fulfilled_idx + 1)



		def __UpperCamelCase				(				self       ,   A_=False )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  PhrasalConstraint(self.token_ids )

				if stateful:
						UpperCamelCase						=  self.seqlen
						UpperCamelCase						=  self.fulfilled_idx
						UpperCamelCase						=  self.completed

				return new_constraint






class        lowercase      :



		def __init__(				self       ,   A_       ,   A_=True )						->      List[Any]:
				"""simple docstring"""
				UpperCamelCase						=  max([len(A_ ) for one in nested_token_ids] )

				UpperCamelCase						=  {}
				for token_ids in nested_token_ids:
						UpperCamelCase						=  root
						for tidx, token_id in enumerate(A_ ):
								if token_id not in level:
										UpperCamelCase						=  {}

								UpperCamelCase						=  level[token_id]

				if no_subsets and self.has_subsets(A_       ,   A_ ):
						raise ValueError(
						    'Each list in `nested_token_ids` can\'t be a complete subset of another list, but is'
						    F''' {nested_token_ids}.''' )

				UpperCamelCase						=  root



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.trie

				for current_token in current_seq:
						UpperCamelCase						=  start[current_token]

				UpperCamelCase						=  list(start.keys() )

				return next_tokens



		def __UpperCamelCase				(				self       ,   A_ )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.next_tokens(A_ )

				return len(A_ ) == 0



		def __UpperCamelCase				(				self       ,   A_ )						->      List[str]:
				"""simple docstring"""
				UpperCamelCase						=  list(root.values() )
				if len(A_ ) == 0:
						return 1
				else:
						return sum([self.count_leaves(A_ ) for nn in next_nodes] )



		def __UpperCamelCase				(				self       ,   A_       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  self.count_leaves(A_ )
				return len(A_ ) != leaf_count






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_ )						->      str:
				"""simple docstring"""
				super(A_       ,   self ).__init__()

				if not isinstance(A_       ,   A_ ) or len(A_ ) == 0:
						raise ValueError(F'''`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}.''' )
				if any(not isinstance(A_       ,   A_ ) for token_ids in nested_token_ids ):
						raise ValueError(F'''`nested_token_ids` has to be a list of lists, but is {nested_token_ids}.''' )
				if any(
				    any((not isinstance(A_       ,   A_ ) or token_id < 0) for token_id in token_ids )
				    for token_ids in nested_token_ids ):
						raise ValueError(
						    F'''Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}.''' )

				UpperCamelCase						=  DisjunctiveTrie(A_ )
				UpperCamelCase						=  nested_token_ids

				UpperCamelCase						=  self.trie.max_height
				UpperCamelCase						=  []
				UpperCamelCase						=  False



		def __UpperCamelCase				(				self )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.trie.next_tokens(self.current_seq )

				if len(A_ ) == 0:
						return None
				else:
						return token_list



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(A_ )}''' )

				UpperCamelCase						=  self.trie.next_tokens(self.current_seq )

				return token_id in next_tokens



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(A_ )}''' )

				UpperCamelCase						=  False
				UpperCamelCase						=  False
				UpperCamelCase						=  False

				if self.does_advance(A_ ):
						self.current_seq.append(A_ )
						UpperCamelCase						=  True
				else:
						UpperCamelCase						=  True
						self.reset()

				UpperCamelCase						=  self.trie.reached_leaf(self.current_seq )
				UpperCamelCase						=  completed

				return stepped, completed, reset



		def __UpperCamelCase				(				self )						->      str:
				"""simple docstring"""
				UpperCamelCase						=  False
				UpperCamelCase						=  []



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				if self.completed:
						# since this can be completed without reaching max height
						return 0
				else:
						return self.seqlen - len(self.current_seq )



		def __UpperCamelCase				(				self       ,   A_=False )						->      int:
				"""simple docstring"""
				UpperCamelCase						=  DisjunctiveConstraint(self.token_ids )

				if stateful:
						UpperCamelCase						=  self.seqlen
						UpperCamelCase						=  self.current_seq
						UpperCamelCase						=  self.completed

				return new_constraint






class        lowercase      :



		def __init__(				self       ,   A_ )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  constraints

				# max # of steps required to fulfill a given constraint
				UpperCamelCase						=  max([c.seqlen for c in constraints] )
				UpperCamelCase						=  len(A_ )
				UpperCamelCase						=  False

				self.init_state()



		def __UpperCamelCase				(				self )						->      List[str]:
				"""simple docstring"""
				UpperCamelCase						=  []
				UpperCamelCase						=  None
				UpperCamelCase						=  [constraint.copy(stateful=A_ ) for constraint in self.constraints]



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  0
				if self.inprogress_constraint:
						# extra points for having a constraint mid-fulfilled
						add += self.max_seqlen - self.inprogress_constraint.remaining()

				return (len(self.complete_constraints ) * self.max_seqlen) + add



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  []
				if self.inprogress_constraint is None:
						for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
								UpperCamelCase						=  constraint.advance()
								if isinstance(A_       ,   A_ ):
										token_list.append(A_ )
								elif isinstance(A_       ,   A_ ):
										token_list.extend(A_ )
				else:
						UpperCamelCase						=  self.inprogress_constraint.advance()
						if isinstance(A_       ,   A_ ):
								token_list.append(A_ )
						elif isinstance(A_       ,   A_ ):
								token_list.extend(A_ )

				if len(A_ ) == 0:
						return None
				else:
						return token_list



		def __UpperCamelCase				(				self       ,   A_ )						->      Any:
				"""simple docstring"""
				self.init_state()

				if token_ids is not None:
						for token in token_ids:
								# completes or steps **one** constraint
								UpperCamelCase			,						UpperCamelCase						=  self.add(A_ )

								# the entire list of constraints are fulfilled
								if self.completed:
										break



		def __UpperCamelCase				(				self       ,   A_ )						->      int:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` should be an `int`, but is `{token_id}`.''' )

				UpperCamelCase			,						UpperCamelCase						=  False, False

				if self.completed:
						UpperCamelCase						=  True
						UpperCamelCase						=  False
						return complete, stepped

				if self.inprogress_constraint is not None:
						# In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current
						# job, simply update the state

						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase						=  self.inprogress_constraint.update(A_ )
						if reset:
								# 1. If the next token breaks the progress, then we must restart.
								#     e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books".

								#     But that doesn't mean we self.init_state(), since we only reset the state for this particular
								#     constraint, not the full list of constraints.

								self.pending_constraints.append(self.inprogress_constraint.copy(stateful=A_ ) )
								UpperCamelCase						=  None

						if complete:
								# 2. If the next token completes the constraint, move it to completed list, set
								#     inprogress to None. If there are no pending constraints either, then this full list of constraints
								#     is complete.

								self.complete_constraints.append(self.inprogress_constraint )
								UpperCamelCase						=  None

								if len(self.pending_constraints ) == 0:
										# we're done!
										UpperCamelCase						=  True

				else:
						# Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list
						# of constraints?

						for cidx, pending_constraint in enumerate(self.pending_constraints ):
								if pending_constraint.does_advance(A_ ):
										UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase						=  pending_constraint.update(A_ )

										if not stepped:
												raise Exception(
												    '`constraint.update(token_id)` is not yielding incremental progress, '
												    'even though `constraint.does_advance(token_id)` is true.' )

										if complete:
												self.complete_constraints.append(A_ )
												UpperCamelCase						=  None

										if not complete and stepped:
												UpperCamelCase						=  pending_constraint

										if complete or stepped:
												# If we made any progress at all, then it's at least not a "pending constraint".

												UpperCamelCase						=  (
												    self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :]
												)

												if len(self.pending_constraints ) == 0 and self.inprogress_constraint is None:
														# If there's no longer any pending after this and no inprogress either, then we must be
														# complete.

														UpperCamelCase						=  True

												break  # prevent accidentally stepping through multiple constraints with just one token.

				return complete, stepped



		def __UpperCamelCase				(				self       ,   A_=True )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  ConstraintListState(self.constraints )  # we actually never though self.constraints objects
				# throughout this process. So it's at initialization state.

				if stateful:
						UpperCamelCase						=  [
						    constraint.copy(stateful=A_ ) for constraint in self.complete_constraints
						]
						if self.inprogress_constraint is not None:
								UpperCamelCase						=  self.inprogress_constraint.copy(stateful=A_ )
						UpperCamelCase						=  [constraint.copy() for constraint in self.pending_constraints]

				return new_state
