





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 :  Union[str, Any]										=  logging.getLogger(__name__)



def 	A_       (  _lowerCAmelCase=2							,     _lowerCAmelCase=3							,     _lowerCAmelCase=16							,     _lowerCAmelCase = 10							,     _lowerCAmelCase = 2  )		->				Dict:

 def get_dataset(_lowerCAmelCase  ):
  UpperCamelCase			:      str 			=     torch.randn(batch_size * n_batches							,     1  )
  return TensorDataset(_lowerCAmelCase							,     a * x + b + 0.1 * torch.randn(batch_size * n_batches							,     1  )  )

 UpperCamelCase			:      Tuple 			=     get_dataset(_lowerCAmelCase  )
 UpperCamelCase			:      Union[str, Any] 			=     get_dataset(_lowerCAmelCase  )
 UpperCamelCase			:      int 			=     DataLoader(_lowerCAmelCase							,     shuffle=_lowerCAmelCase							,     batch_size=_lowerCAmelCase							,     num_workers=4  )
 UpperCamelCase			:      Union[str, Any] 			=     DataLoader(_lowerCAmelCase							,     shuffle=_lowerCAmelCase							,     batch_size=_lowerCAmelCase							,     num_workers=4  )
 return (train_dataloader, valid_dataloader)



def 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase=None  )		->				Optional[Any]:
 UpperCamelCase			:      List[str] 			=     []
 for epoch in range(_lowerCAmelCase  ):
  # Train quickly
  model.train()
  for batch in dataloader:
   UpperCamelCase			,					UpperCamelCase			:      str 			=     batch
   UpperCamelCase			:      int 			=     model(_lowerCAmelCase  )
   UpperCamelCase			:      Tuple 			=     torch.nn.functional.mse_loss(_lowerCAmelCase							,     _lowerCAmelCase  )
   accelerator.backward(_lowerCAmelCase  )
   optimizer.step()
   optimizer.zero_grad()
  rands.append(random.random()  )  # Introduce some randomness
  if scheduler is not None:
   scheduler.step()
 return rands



class       A__     ( nn.Module    ):





      def __init__(				self     ):
       '''simple docstring'''




       super().__init__()
       UpperCamelCase			:      int 			=     nn.Parameter(torch.randn(1     )     )
       UpperCamelCase			:      int 			=     nn.Parameter(torch.randn(1     )     )





      def  __UpperCamelCase(				self ,							A_     ):
       '''simple docstring'''




       return x * self.a + self.b



class       A__     ( unittest.TestCase    ):





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




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

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





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




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

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

        UpperCamelCase			:      Tuple 			=     train(2 ,							A_ ,							A_ ,							A_ ,							A_     )
        # Save everything
        UpperCamelCase			:      Optional[int] 			=     os.path.join(A_ ,							"checkpoint"     )
        accelerator.save_state(A_     )

        # Load everything back in and make sure all states work
        accelerator.load_state(A_     )
        test_rands += train(1 ,							A_ ,							A_ ,							A_ ,							A_     )
        ((UpperCamelCase)			,					(UpperCamelCase))			:      Any 			=     model.a.item(), model.b.item()
        UpperCamelCase			:      Any 			=     optimizer.state_dict()
        self.assertEqual(A_ ,							A_     )
        self.assertEqual(A_ ,							A_     )
        self.assertEqual(A_ ,							A_     )
        self.assertEqual(A_ ,							A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       with tempfile.TemporaryDirectory() as tmpdir:
        set_seed(42     )
        UpperCamelCase			:      Dict 			=     DummyModel()
        UpperCamelCase			:      str 			=     torch.optim.Adam(params=model.parameters() ,							lr=1e-3     )
        UpperCamelCase			,					UpperCamelCase			:      int 			=     dummy_dataloaders()
        UpperCamelCase			:      Dict 			=     ProjectConfiguration(automatic_checkpoint_naming=A_     )

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

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

        UpperCamelCase			:      str 			=     train(2 ,							A_ ,							A_ ,							A_ ,							A_     )
        # Save everything
        accelerator.save_state()

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





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      str 			=     torch.tensor([1, 2, 3]     )
       UpperCamelCase			:      Any 			=     torch.tensor([2, 3, 4]     )
       UpperCamelCase			:      Dict 			=     DummyModel()
       UpperCamelCase			:      Union[str, Any] 			=     torch.optim.Adam(net.parameters()     )
       UpperCamelCase			:      List[str] 			=     Accelerator()
       with self.assertRaises(A_     ) as ve:
        accelerator.register_for_checkpointing(A_ ,							A_ ,							A_ ,							A_     )
       UpperCamelCase			:      List[str] 			=     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  __UpperCamelCase(				self     ):
       '''simple docstring'''




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

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





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       with tempfile.TemporaryDirectory() as tmpdir:
        set_seed(42     )
        UpperCamelCase			:      List[str] 			=     DummyModel()
        UpperCamelCase			:      int 			=     ProjectConfiguration(automatic_checkpoint_naming=A_ ,							total_limit=2     )
        # Train baseline
        UpperCamelCase			:      Optional[int] 			=     Accelerator(project_dir=A_ ,							project_config=A_     )
        UpperCamelCase			:      Tuple 			=     accelerator.prepare(A_     )
        # Save 3 states:
        for _ in range(11     ):
         accelerator.save_state()
        self.assertTrue(not os.path.exists(os.path.join(A_ ,							"checkpoints" ,							"checkpoint_0"     )     )     )
        self.assertTrue(os.path.exists(os.path.join(A_ ,							"checkpoints" ,							"checkpoint_9"     )     )     )
        self.assertTrue(os.path.exists(os.path.join(A_ ,							"checkpoints" ,							"checkpoint_10"     )     )     )





      @require_cuda
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      List[Any] 			=     ["torchrun", F"""--nproc_per_node={torch.cuda.device_count()}""", inspect.getfile(self.__class__     )]
       execute_subprocess_async(A_ ,							env=os.environ.copy()     )


if __name__ == "__main__":
    __lowerCamelCase :  str										=  """/tmp/accelerate/state_checkpointing"""
    __lowerCamelCase :  int										=  DummyModel()
    __lowerCamelCase :  Optional[Any]										=  torch.optim.Adam(params=model.parameters(), lr=1E-3)
    __lowerCamelCase :  Any										=  torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9_9)
    __lowerCamelCase						,			__lowerCamelCase :  Dict										=  dummy_dataloaders()
    __lowerCamelCase :  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 :  Any										=  accelerator.prepare(
        model, optimizer, train_dataloader, valid_dataloader, scheduler
    )
    __lowerCamelCase						,			__lowerCamelCase :  List[Any]										=  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 :  int										=  group["""params"""][0].device
        break
    assert param_device.type == accelerator.device.type
    __lowerCamelCase :  List[str]										=  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 :  Optional[Any]										=  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 :  Optional[int]										=  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()







from __future__ import annotations

__lowerCamelCase :  Optional[int]										=  """Muhammad Umer Farooq"""
__lowerCamelCase :  Tuple										=  """MIT"""
__lowerCamelCase :  Optional[int]										=  """1.0.0"""
__lowerCamelCase :  int										=  """Muhammad Umer Farooq"""
__lowerCamelCase :  Optional[int]										=  """contact@muhammadumerfarooq.me"""
__lowerCamelCase :  Dict										=  """Alpha"""

import re
from html.parser import HTMLParser
from urllib import parse

import requests



class       A__     ( __snake_case    ):





      def __init__(				self ,							A_     ):
       '''simple docstring'''




       super().__init__()
       UpperCamelCase			:      list[str] 			=     []
       UpperCamelCase			:      str 			=     domain





      def  __UpperCamelCase(				self ,							A_ ,							A_     ):
       '''simple docstring'''




       if tag == "a":
        # Check the list of defined attributes.
        for name, value in attrs:
         # If href is defined, and not empty nor # print it.
         if name == "href" and value != "#" and value != "":
          # If not already in urls.
          if value not in self.urls:
           UpperCamelCase			:      Any 			=     parse.urljoin(self.domain ,							A_     )
           self.urls.append(A_     )



def 	A_       (  _lowerCAmelCase  )		->				str:
 return ".".join(get_sub_domain_name(_lowerCAmelCase  ).split("."  )[-2:]  )



def 	A_       (  _lowerCAmelCase  )		->				str:
 return parse.urlparse(_lowerCAmelCase  ).netloc



def 	A_       (  _lowerCAmelCase = "https://github.com"  )		->				list[str]:
 UpperCamelCase			:      int 			=     get_domain_name(_lowerCAmelCase  )

 # Initialize the parser
 UpperCamelCase			:      str 			=     Parser(_lowerCAmelCase  )

 try:
  # Open URL
  UpperCamelCase			:      int 			=     requests.get(_lowerCAmelCase  )

  # pass the raw HTML to the parser to get links
  parser.feed(r.text  )

  # Get links and loop through
  UpperCamelCase			:      Optional[Any] 			=     set()
  for link in parser.urls:
   # open URL.
   # read = requests.get(link)
   try:
    UpperCamelCase			:      Optional[Any] 			=     requests.get(_lowerCAmelCase  )
    # Get the valid email.
    UpperCamelCase			:      Optional[int] 			=     re.findall("[a-zA-Z0-9]+@" + domain							,     read.text  )
    # If not in list then append it.
    for email in emails:
     valid_emails.add(_lowerCAmelCase  )
   except ValueError:
    pass
 except ValueError:
  raise SystemExit(1  )

 # Finally return a sorted list of email addresses with no duplicates.
 return sorted(_lowerCAmelCase  )


if __name__ == "__main__":
    __lowerCamelCase :  Tuple										=  emails_from_url("""https://github.com""")
    print(f"""{len(emails)} emails found:""")
    print("""\n""".join(sorted(emails)))
