


"""simple docstring"""

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





class    lowerCamelCase__						(		_a    ):
      def __init__(		self	:	Optional[Any]							):






       '''simple docstring'''


       self.test()
      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple							):






       '''simple docstring'''


       __lowercase									=	0
       __lowercase									=	False
       while not completed:
        if counter == 1:
         self.reset()
        __lowercase									=	self.advance()
        if not self.does_advance(A_							):
         raise Exception(
             """Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true."""							)

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

        if counter > 1_0_0_0_0:
         raise Exception("""update() does not fulfill the constraint."""							)

       if self.remaining() != 0:
        raise Exception("""Custom Constraint is not defined correctly."""							)
      @abstractmethod
      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple							):






       '''simple docstring'''


       raise NotImplementedError(
           F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''							)
      @abstractmethod
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]      ,       A_	:	int							):






       '''simple docstring'''


       raise NotImplementedError(
           F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''							)
      @abstractmethod
      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]      ,       A_	:	int							):






       '''simple docstring'''


       raise NotImplementedError(
           F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''							)
      @abstractmethod
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       raise NotImplementedError(
           F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''							)
      @abstractmethod
      def SCREAMING_SNAKE_CASE_	(		self	:	Union[str, Any]							):






       '''simple docstring'''


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






      @abstractmethod
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]      ,       A_	:	List[str]=False							):






       '''simple docstring'''


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






class    lowerCamelCase__						(		_a    ):
      def __init__(		self	:	Optional[int]      ,       A_	:	List[int]							):






       '''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}.'''							)

       __lowercase									=	token_ids

       __lowercase									=	len(self.token_ids							)
       __lowercase									=	-1  # the index of the currently fulfilled step
       __lowercase									=	False
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       if self.completed:
        return None
       return self.token_ids[self.fulfilled_idx + 1]
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict      ,       A_	:	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 SCREAMING_SNAKE_CASE_	(		self	:	Union[str, Any]      ,       A_	:	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_							)}'''							)

       __lowercase									=	False
       __lowercase									=	False
       __lowercase									=	False

       if self.does_advance(A_							):
        self.fulfilled_idx += 1
        __lowercase									=	True
        if self.fulfilled_idx == (self.seqlen - 1):
         __lowercase									=	True
        __lowercase									=	completed
       else:
        # failed to make progress.
        __lowercase									=	True
        self.reset()
       return stepped, completed, reset
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]							):






       '''simple docstring'''


       __lowercase									=	False
       __lowercase									=	0
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]							):






       '''simple docstring'''


       return self.seqlen - (self.fulfilled_idx + 1)






      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]      ,       A_	:	List[str]=False							):






       '''simple docstring'''


       __lowercase									=	PhrasalConstraint(self.token_ids							)

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

       return new_constraint






class    lowerCamelCase__						:
      def __init__(		self	:	Tuple      ,       A_	:	List[List[int]]      ,       A_	:	int=True							):






       '''simple docstring'''


       __lowercase									=	max([len(A_							) for one in nested_token_ids]							)

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

         __lowercase									=	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}.'''							)

       __lowercase									=	root
      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]      ,       A_	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	self.trie

       for current_token in current_seq:
        __lowercase									=	start[current_token]

       __lowercase									=	list(start.keys()							)

       return next_tokens
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]      ,       A_	:	List[str]							):






       '''simple docstring'''


       __lowercase									=	self.next_tokens(A_							)

       return len(A_							) == 0
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict      ,       A_	:	Dict							):






       '''simple docstring'''


       __lowercase									=	list(root.values()							)
       if len(A_							) == 0:
        return 1
       else:
        return sum([self.count_leaves(A_							) for nn in next_nodes]							)






      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple      ,       A_	:	Optional[Any]      ,       A_	:	Tuple							):






       '''simple docstring'''


       __lowercase									=	self.count_leaves(A_							)
       return len(A_							) != leaf_count






class    lowerCamelCase__						(		_a    ):
      def __init__(		self	:	Tuple      ,       A_	:	List[List[int]]							):






       '''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}.'''							)

       __lowercase									=	DisjunctiveTrie(A_							)
       __lowercase									=	nested_token_ids

       __lowercase									=	self.trie.max_height
       __lowercase									=	[]
       __lowercase									=	False
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


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

       if len(A_							) == 0:
        return None
       else:
        return token_list
      def SCREAMING_SNAKE_CASE_	(		self	:	int      ,       A_	:	int							):






       '''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_							)}'''							)

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

       return token_id in next_tokens
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]      ,       A_	:	int							):






       '''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_							)}'''							)

       __lowercase									=	False
       __lowercase									=	False
       __lowercase									=	False

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

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

       return stepped, completed, reset
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	False
       __lowercase									=	[]
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''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 SCREAMING_SNAKE_CASE_	(		self	:	List[str]      ,       A_	:	Optional[Any]=False							):






       '''simple docstring'''


       __lowercase									=	DisjunctiveConstraint(self.token_ids							)

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

       return new_constraint






class    lowerCamelCase__						:
      def __init__(		self	:	List[Any]      ,       A_	:	List[Constraint]							):






       '''simple docstring'''


       __lowercase									=	constraints

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

       self.init_state()
      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple							):






       '''simple docstring'''


       __lowercase									=	[]
       __lowercase									=	None
       __lowercase									=	[constraint.copy(stateful=A_							) for constraint in self.constraints]
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	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 SCREAMING_SNAKE_CASE_	(		self	:	str							):






       '''simple docstring'''


       __lowercase									=	[]
       if self.inprogress_constraint is None:
        for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
         __lowercase									=	constraint.advance()
         if isinstance(A_      ,       A_							):
          token_list.append(A_							)
         elif isinstance(A_      ,       A_							):
          token_list.extend(A_							)
       else:
        __lowercase									=	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 SCREAMING_SNAKE_CASE_	(		self	:	Any      ,       A_	:	Optional[List[int]]							):






       '''simple docstring'''


       self.init_state()

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

         # the entire list of constraints are fulfilled
         if self.completed:
          break
      def SCREAMING_SNAKE_CASE_	(		self	:	Union[str, Any]      ,       A_	:	int							):






       '''simple docstring'''


       if not isinstance(A_      ,       A_							):
        raise ValueError(F'''`token_id` should be an `int`, but is `{token_id}`.'''							)

       __lowercase							,		__lowercase									=	False, False

       if self.completed:
        __lowercase									=	True
        __lowercase									=	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

        __lowercase							,		__lowercase							,		__lowercase									=	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_							)							)
         __lowercase									=	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							)
         __lowercase									=	None

         if len(self.pending_constraints							) == 0:
          # we're done!
          __lowercase									=	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_							):
          __lowercase							,		__lowercase							,		__lowercase									=	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_							)
           __lowercase									=	None

          if not complete and stepped:
           __lowercase									=	pending_constraint

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

           __lowercase									=	(
               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.

            __lowercase									=	True

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

       return complete, stepped






      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]      ,       A_	:	int=True							):






       '''simple docstring'''


       __lowercase									=	ConstraintListState(self.constraints							)  # we actually never though self.constraints objects
       # throughout this process. So it's at initialization state.

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

       return new_state




"""simple docstring"""

import inspect
import unittest





class    lowerCamelCase__						(		unittest.TestCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	int							):






       '''simple docstring'''


       try:
        import diffusers  # noqa: F401
       except ImportError:
        assert False






      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       import diffusers
       from diffusers.dependency_versions_table import deps

       __lowercase									=	inspect.getmembers(A_      ,       inspect.isclass							)

       for cls_name, cls_module in all_classes:
        if "dummy_" in cls_module.__module__:
         for backend in cls_module._backends:
          if backend == "k_diffusion":
           __lowercase									=	"""k-diffusion"""
          elif backend == "invisible_watermark":
           __lowercase									=	"""invisible-watermark"""
          assert backend in deps, F'''{backend} is not in the deps table!'''
