

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



class       __snake_case  (    __a       ):

  def __init__(					self):
      self.test()

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						0
      SCREAMING_SNAKE_CASE_									=						False
      while not completed:
          if counter == 1:
              self.reset()
          SCREAMING_SNAKE_CASE_									=						self.advance()
          if not self.does_advance(a_):
              raise Exception(
                  'Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.')

          SCREAMING_SNAKE_CASE_									=						self.update(a_)
          counter += 1

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

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

  @abstractmethod
  def 						lowerCAmelCase__     (					self):
      raise NotImplementedError(
          f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called.""")

  @abstractmethod
  def 						lowerCAmelCase__     (					self						,		_A):
      raise NotImplementedError(
          f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called.""")

  @abstractmethod
  def 						lowerCAmelCase__     (					self						,		_A):
      raise NotImplementedError(
          f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called.""")

  @abstractmethod
  def 						lowerCAmelCase__     (					self):
      raise NotImplementedError(
          f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called.""")

  @abstractmethod
  def 						lowerCAmelCase__     (					self):
      raise NotImplementedError(
          f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called.""")


  @abstractmethod
  def 						lowerCAmelCase__     (					self						,		_A=False):
      raise NotImplementedError(
          f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called.""")






class       __snake_case  (    __a       ):

  def __init__(					self						,		_A):
      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}.""")

      SCREAMING_SNAKE_CASE_									=						token_ids

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

  def 						lowerCAmelCase__     (					self):



      if self.completed:
          return None
      return self.token_ids[self.fulfilled_idx + 1]

  def 						lowerCAmelCase__     (					self						,		_A):



      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 						lowerCAmelCase__     (					self						,		_A):



      if not isinstance(a_						,		a_):
          raise ValueError(f"""`token_id` has to be an `int`, but is {token_id} of type {type(a_)}""")

      SCREAMING_SNAKE_CASE_									=						False
      SCREAMING_SNAKE_CASE_									=						False
      SCREAMING_SNAKE_CASE_									=						False

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

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						False
      SCREAMING_SNAKE_CASE_									=						0

  def 						lowerCAmelCase__     (					self):
      return self.seqlen - (self.fulfilled_idx + 1)


  def 						lowerCAmelCase__     (					self						,		_A=False):
      SCREAMING_SNAKE_CASE_									=						PhrasalConstraint(self.token_ids)

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

      return new_constraint






class       __snake_case  :

  def __init__(					self						,		_A						,		_A=True):
      SCREAMING_SNAKE_CASE_									=						max([len(a_) for one in nested_token_ids])

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

              SCREAMING_SNAKE_CASE_									=						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}.""")

      SCREAMING_SNAKE_CASE_									=						root

  def 						lowerCAmelCase__     (					self						,		_A):
      SCREAMING_SNAKE_CASE_									=						self.trie

      for current_token in current_seq:
          SCREAMING_SNAKE_CASE_									=						start[current_token]

      SCREAMING_SNAKE_CASE_									=						list(start.keys())

      return next_tokens

  def 						lowerCAmelCase__     (					self						,		_A):
      SCREAMING_SNAKE_CASE_									=						self.next_tokens(a_)

      return len(a_) == 0

  def 						lowerCAmelCase__     (					self						,		_A):
      SCREAMING_SNAKE_CASE_									=						list(root.values())
      if len(a_) == 0:
          return 1
      else:
          return sum([self.count_leaves(a_) for nn in next_nodes])


  def 						lowerCAmelCase__     (					self						,		_A						,		_A):
      SCREAMING_SNAKE_CASE_									=						self.count_leaves(a_)
      return len(a_) != leaf_count






class       __snake_case  (    __a       ):

  def __init__(					self						,		_A):
      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}.""")

      SCREAMING_SNAKE_CASE_									=						DisjunctiveTrie(a_)
      SCREAMING_SNAKE_CASE_									=						nested_token_ids

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

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.trie.next_tokens(self.current_seq)

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

  def 						lowerCAmelCase__     (					self						,		_A):



      if not isinstance(a_						,		a_):
          raise ValueError(f"""`token_id` is supposed to be type `int`, but is {token_id} of type {type(a_)}""")

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

      return token_id in next_tokens

  def 						lowerCAmelCase__     (					self						,		_A):



      if not isinstance(a_						,		a_):
          raise ValueError(f"""`token_id` is supposed to be type `int`, but is {token_id} of type {type(a_)}""")

      SCREAMING_SNAKE_CASE_									=						False
      SCREAMING_SNAKE_CASE_									=						False
      SCREAMING_SNAKE_CASE_									=						False

      if self.does_advance(a_):
          self.current_seq.append(a_)
          SCREAMING_SNAKE_CASE_									=						True
      else:
          SCREAMING_SNAKE_CASE_									=						True
          self.reset()

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

      return stepped, completed, reset

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						False
      SCREAMING_SNAKE_CASE_									=						[]

  def 						lowerCAmelCase__     (					self):



      if self.completed:
          # since this can be completed without reaching max height
          return 0
      else:
          return self.seqlen - len(self.current_seq)


  def 						lowerCAmelCase__     (					self						,		_A=False):
      SCREAMING_SNAKE_CASE_									=						DisjunctiveConstraint(self.token_ids)

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

      return new_constraint






class       __snake_case  :

  def __init__(					self						,		_A):
      SCREAMING_SNAKE_CASE_									=						constraints

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

      self.init_state()

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						[]
      SCREAMING_SNAKE_CASE_									=						None
      SCREAMING_SNAKE_CASE_									=						[constraint.copy(stateful=a_) for constraint in self.constraints]

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						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 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						[]
      if self.inprogress_constraint is None:
          for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
              SCREAMING_SNAKE_CASE_									=						constraint.advance()
              if isinstance(a_						,		a_):
                  token_list.append(a_)
              elif isinstance(a_						,		a_):
                  token_list.extend(a_)
      else:
          SCREAMING_SNAKE_CASE_									=						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 						lowerCAmelCase__     (					self						,		_A):
      self.init_state()

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

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

  def 						lowerCAmelCase__     (					self						,		_A):



      if not isinstance(a_						,		a_):
          raise ValueError(f"""`token_id` should be an `int`, but is `{token_id}`.""")

      SCREAMING_SNAKE_CASE_									=						False, False

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

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

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

                  if not complete and stepped:
                      SCREAMING_SNAKE_CASE_									=						pending_constraint

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

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

                          SCREAMING_SNAKE_CASE_									=						True

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

      return complete, stepped


  def 						lowerCAmelCase__     (					self						,		_A=True):
      SCREAMING_SNAKE_CASE_									=						ConstraintListState(self.constraints)  # we actually never though self.constraints objects
      # throughout this process. So it's at initialization state.

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

      return new_state





import inspect
import os
import unittest
from pathlib import Path

import torch

import accelerate
from accelerate.test_utils import execute_subprocess_async
from accelerate.test_utils.testing import run_command



class       __snake_case  (    unittest.TestCase       ):
  __lowerCAmelCase     :	Dict						 =  inspect.getfile(accelerate.test_utils       )
  __lowerCAmelCase     :	Optional[Any]						 =  os.path.sep.join(mod_file.split(os.path.sep       )[:-1] + ['scripts', 'test_cli.py']       )

  __lowerCAmelCase     :	Tuple						 =  ['accelerate', 'launch']
  __lowerCAmelCase     :	Union[str, Any]						 =  Path.home() / '.cache/huggingface/accelerate'
  __lowerCAmelCase     :	List[str]						 =  'default_config.yaml'
  __lowerCAmelCase     :	List[Any]						 =  config_folder / config_file
  __lowerCAmelCase     :	str						 =  config_folder / '_default_config.yaml'

  __lowerCAmelCase     :	Optional[int]						 =  Path('tests/test_configs'       )

  @classmethod
  def 						lowerCAmelCase__     (					cls):
      if cls.config_path.is_file():
          cls.config_path.rename(cls.changed_path)

  @classmethod
  def 						lowerCAmelCase__     (					cls):
      if cls.changed_path.is_file():
          cls.changed_path.rename(cls.config_path)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.base_cmd
      if torch.cuda.is_available() and (torch.cuda.device_count() > 1):
          cmd += ["--multi_gpu"]
      execute_subprocess_async(cmd + [self.test_file_path]						,		env=os.environ.copy())

  def 						lowerCAmelCase__     (					self):
      for config in sorted(self.test_config_path.glob('**/*.yaml')):
          with self.subTest(config_file=_A):
              execute_subprocess_async(
                  self.base_cmd + ['--config_file', str(_A), self.test_file_path]						,		env=os.environ.copy())


  def 						lowerCAmelCase__     (					self):
      execute_subprocess_async(['accelerate', 'test']						,		env=os.environ.copy())






class       __snake_case  (    unittest.TestCase       ):
  __lowerCAmelCase     :	Optional[Any]						 =  'test-tpu'
  __lowerCAmelCase     :	str						 =  'us-central1-a'
  __lowerCAmelCase     :	Union[str, Any]						 =  'ls'
  __lowerCAmelCase     :	Union[str, Any]						 =  ['accelerate', 'tpu-config']
  __lowerCAmelCase     :	Union[str, Any]						 =  'cd /usr/share'
  __lowerCAmelCase     :	List[Any]						 =  'tests/test_samples/test_command_file.sh'
  __lowerCAmelCase     :	Dict						 =  'Running gcloud compute tpus tpu-vm ssh'

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + ['--command', self.command, '--tpu_zone', self.tpu_zone, '--tpu_name', self.tpu_name, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/0_12_0.yaml',
              '--command',
              self.command,
              '--tpu_zone',
              self.tpu_zone,
              '--tpu_name',
              self.tpu_name,
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--debug']						,		return_stdout=_A)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--command', self.command, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/latest.yaml',
              '--command',
              self.command,
              '--command',
              'echo "Hello World"',
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls; echo \"Hello World\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + ['--config_file', 'tests/test_configs/latest.yaml', '--command_file', self.command_file, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/0_12_0.yaml',
              '--command_file',
              self.command_file,
              '--tpu_zone',
              self.tpu_zone,
              '--tpu_name',
              self.tpu_name,
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--install_accelerate', '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate -U; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/latest.yaml',
              '--install_accelerate',
              '--accelerate_version',
              '12.0.0',
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate==12.0.0; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)


