
import inspect
import logging
import os
import random
import shutil
import tempfile
import unittest

import pytest
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_cuda
from accelerate.utils import ProjectConfiguration, set_seed


lowerCamelCase__    :		Optional[Any]          =   logging.getLogger(__name__)

def 		UpperCamelCase	(      lowercase_=2 , lowercase_=3 , lowercase_=16 , lowercase_ = 10 , lowercase_ = 2   )		->	Tuple:


  '''simple docstring'''






  def get_dataset(lowercase_   ):
    lowercase__  :   List[str]    		=      torch.randn(batch_size * n_batches , 1   )
    return TensorDataset(lowercase_ , a * x + b + 0.1 * torch.randn(batch_size * n_batches , 1   )   )

  lowercase__  :   str    		=      get_dataset(lowercase_   )
  lowercase__  :   Union[str, Any]    		=      get_dataset(lowercase_   )
  lowercase__  :   Union[str, Any]    		=      DataLoader(lowercase_ , shuffle=lowercase_ , batch_size=lowercase_ , num_workers=4   )
  lowercase__  :   Dict    		=      DataLoader(lowercase_ , shuffle=lowercase_ , batch_size=lowercase_ , num_workers=4   )
  return (train_dataloader, valid_dataloader)

def 		UpperCamelCase	(      lowercase_ , lowercase_ , lowercase_ , lowercase_ , lowercase_ , lowercase_=None   )		->	Dict:


  '''simple docstring'''






  lowercase__  :   Dict    		=      []
  for epoch in range(lowercase_   ):
    # Train quickly
    model.train()
    for batch in dataloader:
      lowercase__		,    lowercase__  :   Union[str, Any]    		=      batch
      lowercase__  :   Dict    		=      model(lowercase_   )
      lowercase__  :   Optional[Any]    		=      torch.nn.functional.mse_loss(lowercase_ , lowercase_   )
      accelerator.backward(lowercase_   )
      optimizer.step()
      optimizer.zero_grad()
    rands.append(random.random()   )  # Introduce some randomness
    if scheduler is not None:
      scheduler.step()
  return rands



class       _snake_case   (	nn.Module							):


 def __init__(					self):


   '''simple docstring'''




   super().__init__()
   lowercase__  :   Optional[Any]    		=      nn.Parameter(torch.randn(1))
   lowercase__  :   Dict    		=      nn.Parameter(torch.randn(1))



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   return x * self.a + self.b



class       _snake_case   (	unittest.TestCase							):


 def      lowercase__     (					self):


   '''simple docstring'''




   with tempfile.TemporaryDirectory() as tmpdir:
     set_seed(42)
     lowercase__  :   Union[str, Any]    		=      DummyModel()
     lowercase__  :   Optional[int]    		=      torch.optim.Adam(params=model.parameters()					,  lr=1E-3)
     lowercase__		,    lowercase__  :   int    		=      dummy_dataloaders()
     lowercase__  :   str    		=      ProjectConfiguration(total_limit=1					,  project_dir=SCREAMING_SNAKE_CASE_					,  automatic_checkpoint_naming=SCREAMING_SNAKE_CASE_)
     # Train baseline
     lowercase__  :   Any    		=      Accelerator(project_config=SCREAMING_SNAKE_CASE_)
     lowercase__		,    lowercase__		,    lowercase__		,    lowercase__  :   Tuple    		=      accelerator.prepare(
         SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     # Save initial
     accelerator.save_state()

     # Save second state
     accelerator.save_state()
     self.assertEqual(len(os.listdir(accelerator.project_dir))					,  1)


 def      lowercase__     (					self):


   '''simple docstring'''




   with tempfile.TemporaryDirectory() as tmpdir:
     set_seed(42)
     lowercase__  :   List[Any]    		=      DummyModel()
     lowercase__  :   List[Any]    		=      torch.optim.Adam(params=model.parameters()					,  lr=1E-3)
     lowercase__		,    lowercase__  :   List[str]    		=      dummy_dataloaders()
     # Train baseline
     lowercase__  :   str    		=      Accelerator()
     lowercase__		,    lowercase__		,    lowercase__		,    lowercase__  :   Tuple    		=      accelerator.prepare(
         SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     # Save initial
     lowercase__  :   Optional[int]    		=      os.path.join(SCREAMING_SNAKE_CASE_					,  """initial""")
     accelerator.save_state(SCREAMING_SNAKE_CASE_)
     ((lowercase__)		,    (lowercase__))  :   List[Any]    		=      model.a.item(), model.b.item()
     lowercase__  :   Any    		=      optimizer.state_dict()
     lowercase__  :   Tuple    		=      train(3					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     ((lowercase__)		,    (lowercase__))  :   str    		=      model.a.item(), model.b.item()
     lowercase__  :   Any    		=      optimizer.state_dict()

     # Train partially
     set_seed(42)
     lowercase__  :   List[str]    		=      DummyModel()
     lowercase__  :   Optional[int]    		=      torch.optim.Adam(params=model.parameters()					,  lr=1E-3)
     lowercase__		,    lowercase__  :   Tuple    		=      dummy_dataloaders()
     lowercase__  :   Tuple    		=      Accelerator()
     lowercase__		,    lowercase__		,    lowercase__		,    lowercase__  :   Tuple    		=      accelerator.prepare(
         SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     accelerator.load_state(SCREAMING_SNAKE_CASE_)
     ((lowercase__)		,    (lowercase__))  :   Tuple    		=      model.a.item(), model.b.item()
     lowercase__  :   Dict    		=      optimizer.state_dict()
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)

     lowercase__  :   int    		=      train(2					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     # Save everything
     lowercase__  :   List[Any]    		=      os.path.join(SCREAMING_SNAKE_CASE_					,  """checkpoint""")
     accelerator.save_state(SCREAMING_SNAKE_CASE_)

     # Load everything back in and make sure all states work
     accelerator.load_state(SCREAMING_SNAKE_CASE_)
     test_rands += train(1					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     ((lowercase__)		,    (lowercase__))  :   Optional[int]    		=      model.a.item(), model.b.item()
     lowercase__  :   Union[str, Any]    		=      optimizer.state_dict()
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)


 def      lowercase__     (					self):


   '''simple docstring'''




   with tempfile.TemporaryDirectory() as tmpdir:
     set_seed(42)
     lowercase__  :   Optional[Any]    		=      DummyModel()
     lowercase__  :   Optional[int]    		=      torch.optim.Adam(params=model.parameters()					,  lr=1E-3)
     lowercase__		,    lowercase__  :   int    		=      dummy_dataloaders()
     lowercase__  :   int    		=      ProjectConfiguration(automatic_checkpoint_naming=SCREAMING_SNAKE_CASE_)

     # Train baseline
     lowercase__  :   Union[str, Any]    		=      Accelerator(project_dir=SCREAMING_SNAKE_CASE_					,  project_config=SCREAMING_SNAKE_CASE_)
     lowercase__		,    lowercase__		,    lowercase__		,    lowercase__  :   Optional[int]    		=      accelerator.prepare(
         SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     # Save initial
     accelerator.save_state()
     ((lowercase__)		,    (lowercase__))  :   Optional[Any]    		=      model.a.item(), model.b.item()
     lowercase__  :   int    		=      optimizer.state_dict()
     lowercase__  :   Optional[int]    		=      train(3					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     ((lowercase__)		,    (lowercase__))  :   List[str]    		=      model.a.item(), model.b.item()
     lowercase__  :   Tuple    		=      optimizer.state_dict()

     # Train partially
     set_seed(42)
     lowercase__  :   Any    		=      DummyModel()
     lowercase__  :   List[Any]    		=      torch.optim.Adam(params=model.parameters()					,  lr=1E-3)
     lowercase__		,    lowercase__  :   Tuple    		=      dummy_dataloaders()
     lowercase__  :   Tuple    		=      ProjectConfiguration(iteration=1					,  automatic_checkpoint_naming=SCREAMING_SNAKE_CASE_)
     lowercase__  :   Optional[Any]    		=      Accelerator(project_dir=SCREAMING_SNAKE_CASE_					,  project_config=SCREAMING_SNAKE_CASE_)
     lowercase__		,    lowercase__		,    lowercase__		,    lowercase__  :   Optional[Any]    		=      accelerator.prepare(
         SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     accelerator.load_state(os.path.join(SCREAMING_SNAKE_CASE_					,  """checkpoints"""					,  """checkpoint_0"""))
     ((lowercase__)		,    (lowercase__))  :   Union[str, Any]    		=      model.a.item(), model.b.item()
     lowercase__  :   str    		=      optimizer.state_dict()
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)

     lowercase__  :   Union[str, Any]    		=      train(2					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     # Save everything
     accelerator.save_state()

     # Load everything back in and make sure all states work
     accelerator.load_state(os.path.join(SCREAMING_SNAKE_CASE_					,  """checkpoints"""					,  """checkpoint_1"""))
     test_rands += train(1					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     ((lowercase__)		,    (lowercase__))  :   str    		=      model.a.item(), model.b.item()
     lowercase__  :   List[Any]    		=      optimizer.state_dict()
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)


 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   List[Any]    		=      torch.tensor([1, 2, 3])
   lowercase__  :   Dict    		=      torch.tensor([2, 3, 4])
   lowercase__  :   List[str]    		=      DummyModel()
   lowercase__  :   List[Any]    		=      torch.optim.Adam(net.parameters())
   lowercase__  :   str    		=      Accelerator()
   with self.assertRaises(SCREAMING_SNAKE_CASE_) as ve:
     accelerator.register_for_checkpointing(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
   lowercase__  :   int    		=      str(ve.exception)
   self.assertTrue("""Item at index 0""" in message)
   self.assertTrue("""Item at index 1""" in message)
   self.assertFalse("""Item at index 2""" in message)
   self.assertFalse("""Item at index 3""" in message)


 def      lowercase__     (					self):


   '''simple docstring'''




   with tempfile.TemporaryDirectory() as tmpdir:
     set_seed(42)
     lowercase__  :   Any    		=      DummyModel()
     lowercase__  :   Optional[int]    		=      torch.optim.Adam(params=model.parameters()					,  lr=1E-3)
     lowercase__  :   Dict    		=      torch.optim.lr_scheduler.StepLR(SCREAMING_SNAKE_CASE_					,  step_size=1					,  gamma=0.9_9)
     lowercase__		,    lowercase__  :   Union[str, Any]    		=      dummy_dataloaders()
     lowercase__  :   Any    		=      ProjectConfiguration(automatic_checkpoint_naming=SCREAMING_SNAKE_CASE_)
     # Train baseline
     lowercase__  :   Optional[Any]    		=      Accelerator(project_dir=SCREAMING_SNAKE_CASE_					,  project_config=SCREAMING_SNAKE_CASE_)
     lowercase__		,    lowercase__		,    lowercase__		,    lowercase__		,    lowercase__  :   Any    		=      accelerator.prepare(
         SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     # Save initial
     accelerator.save_state()
     lowercase__  :   Dict    		=      scheduler.state_dict()
     train(3					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
     self.assertNotEqual(SCREAMING_SNAKE_CASE_					,  scheduler.state_dict())

     # Load everything back in and make sure all states work
     accelerator.load_state(os.path.join(SCREAMING_SNAKE_CASE_					,  """checkpoints"""					,  """checkpoint_0"""))
     self.assertEqual(SCREAMING_SNAKE_CASE_					,  scheduler.state_dict())


 def      lowercase__     (					self):


   '''simple docstring'''




   with tempfile.TemporaryDirectory() as tmpdir:
     set_seed(42)
     lowercase__  :   Any    		=      DummyModel()
     lowercase__  :   Dict    		=      ProjectConfiguration(automatic_checkpoint_naming=SCREAMING_SNAKE_CASE_					,  total_limit=2)
     # Train baseline
     lowercase__  :   Any    		=      Accelerator(project_dir=SCREAMING_SNAKE_CASE_					,  project_config=SCREAMING_SNAKE_CASE_)
     lowercase__  :   Dict    		=      accelerator.prepare(SCREAMING_SNAKE_CASE_)
     # Save 3 states:
     for _ in range(11):
       accelerator.save_state()
     self.assertTrue(not os.path.exists(os.path.join(SCREAMING_SNAKE_CASE_					,  """checkpoints"""					,  """checkpoint_0""")))
     self.assertTrue(os.path.exists(os.path.join(SCREAMING_SNAKE_CASE_					,  """checkpoints"""					,  """checkpoint_9""")))
     self.assertTrue(os.path.exists(os.path.join(SCREAMING_SNAKE_CASE_					,  """checkpoints"""					,  """checkpoint_10""")))



 @require_cuda
 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   str    		=      ["""torchrun""", f'--nproc_per_node={torch.cuda.device_count()}', inspect.getfile(self.__class__)]
   execute_subprocess_async(SCREAMING_SNAKE_CASE_					,  env=os.environ.copy())


if __name__ == "__main__":
     lowerCamelCase__    :		Union[str, Any]          =   """/tmp/accelerate/state_checkpointing"""
     lowerCamelCase__    :		Tuple          =   DummyModel()
     lowerCamelCase__    :		int          =   torch.optim.Adam(params=model.parameters(), lr=1E-3)
     lowerCamelCase__    :		str          =   torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
     lowerCamelCase__      ,						lowerCamelCase__    :		Any          =   dummy_dataloaders()
     lowerCamelCase__    :		Optional[Any]          =   ProjectConfiguration(automatic_checkpoint_naming=True)
     # Train baseline
     lowerCamelCase__    :		Any          =   Accelerator(project_dir=savedir, project_config=project_config, mixed_precision="""no""")
     if accelerator.process_index == 0:
          if os.path.exists(savedir):
               shutil.rmtree(savedir)
          os.makedirs(savedir)
     lowerCamelCase__      ,						lowerCamelCase__      ,						lowerCamelCase__      ,						lowerCamelCase__      ,						lowerCamelCase__    :		int          =   accelerator.prepare(
         model, optimizer, train_dataloader, valid_dataloader, scheduler
     )
     lowerCamelCase__      ,						lowerCamelCase__    :		int          =   accelerator.prepare(model, optimizer)
     train(3, model, train_dataloader, optimizer, accelerator, scheduler)
     # Check that the intial optimizer is loaded on the GPU
     for group in optimizer.param_groups:
          lowerCamelCase__    :		List[str]          =   group["""params"""][0].device
          break
     assert param_device.type == accelerator.device.type
     lowerCamelCase__    :		Optional[int]          =   model.cpu()
     accelerator.wait_for_everyone()
     accelerator.save_state()
     accelerator.wait_for_everyone()

     # Check CPU state
     accelerator.load_state(os.path.join(savedir, """checkpoints""", """checkpoint_0"""), map_location="""cpu""")
     for group in optimizer.param_groups:
          lowerCamelCase__    :		Dict          =   group["""params"""][0].device
          break
     assert (
         param_device.type == torch.device("""cpu""").type
     ), f"Loaded optimizer states did not match, expected to be loaded on the CPU but got {param_device}"

     # Check device state
     model.to(accelerator.device)
     accelerator.load_state(os.path.join(savedir, """checkpoints""", """checkpoint_0"""), map_location="""on_device""")
     for group in optimizer.param_groups:
          lowerCamelCase__    :		Dict          =   group["""params"""][0].device
          break
     assert (
         param_device.type == accelerator.device.type
     ), f"Loaded optimizer states did not match, expected to be loaded on {accelerator.device} but got {param_device}"

     # Check error
     with pytest.raises(TypeError, match="""Unsupported optimizer map location passed"""):
          accelerator.load_state(os.path.join(savedir, """checkpoints""", """checkpoint_0"""), map_location="""invalid""")
     accelerator.wait_for_everyone()
     if accelerator.process_index == 0:
          shutil.rmtree(savedir)
     accelerator.wait_for_everyone()








# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

import torch

from ..models.auto import AutoModelForVisualQuestionAnswering, AutoProcessor
from ..utils import requires_backends
from .base import PipelineTool


if TYPE_CHECKING:
     from PIL import Image



class       _snake_case   (	UpperCAmelCase_							):
 __lowerCAmelCase  :     Any        =			'dandelin/vilt-b32-finetuned-vqa'
 __lowerCAmelCase  :     Union[str, Any]        =			(
     'This is a tool that answers a question about an image. It takes an input named `image` which should be the '
     'image containing the information, as well as a `question` which should be the question in English. It '
     'returns a text that is the answer to the question.'
 )
 __lowerCAmelCase  :     str        =			'image_qa'
 __lowerCAmelCase  :     int        =			AutoProcessor
 __lowerCAmelCase  :     Optional[Any]        =			AutoModelForVisualQuestionAnswering

 __lowerCAmelCase  :     str        =			['image', 'text']
 __lowerCAmelCase  :     Optional[int]        =			['text']


 def __init__(					self					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(self					,  ["""vision"""])
   super().__init__(*SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_)


 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   return self.pre_processor(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  return_tensors="""pt""")


 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   with torch.no_grad():
     return self.model(**SCREAMING_SNAKE_CASE_).logits



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__  :   Dict    		=      outputs.argmax(-1).item()
   return self.model.config.idalabel[idx]






