


'''simple docstring'''






import math
import os
from copy import deepcopy

import datasets
import evaluate
import torch
import transformers
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from accelerate import Accelerator
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import is_tpu_available, set_seed


__A					:  Optional[int]											=      "true"







def 				UpperCamelCase_			(    A__					:   Optional[int]							,    A__					:   List[Any]=82							,    A__					:   Optional[int]=16					):
  '''simple docstring'''





  set_seed(42					)
  lowerCAmelCase_			:      Dict			=   RegressionModel()
  lowerCAmelCase_			:      Dict			=   deepcopy(A__					)
  lowerCAmelCase_			:      int			=   RegressionDataset(length=A__					)
  lowerCAmelCase_			:      List[Any]			=   DataLoader(A__							,    batch_size=A__					)
  model.to(accelerator.device					)
  lowerCAmelCase_,							lowerCAmelCase_			:      Tuple			=   accelerator.prepare(A__							,    A__					)
  return model, ddp_model, dataloader







def 				UpperCamelCase_			(    A__					:   Accelerator							,    A__					:   Optional[Any]=False					):
  '''simple docstring'''





  lowerCAmelCase_			:      Optional[int]			=   AutoTokenizer.from_pretrained("""hf-internal-testing/mrpc-bert-base-cased"""					)
  lowerCAmelCase_			:      str			=   load_dataset("""glue"""							,    """mrpc"""							,    split="""validation"""					)

  def tokenize_function(A__					:   Optional[int]					):
    lowerCAmelCase_			:      Any			=   tokenizer(examples["""sentence1"""]							,    examples["""sentence2"""]							,    truncation=A__							,    max_length=A__					)
    return outputs

  with accelerator.main_process_first():
    lowerCAmelCase_			:      Optional[int]			=   dataset.map(
        A__							,    batched=A__							,    remove_columns=["""idx""", """sentence1""", """sentence2"""]							,    )

  lowerCAmelCase_			:      str			=   tokenized_datasets.rename_column("""label"""							,    """labels"""					)

  def collate_fn(A__					:   Optional[int]					):
    if use_longest:
      return tokenizer.pad(A__							,    padding="""longest"""							,    return_tensors="""pt"""					)
    return tokenizer.pad(A__							,    padding="""max_length"""							,    max_length=1_28							,    return_tensors="""pt"""					)

  return DataLoader(A__							,    shuffle=A__							,    collate_fn=A__							,    batch_size=16					)







def 				UpperCamelCase_			(    A__					:   str							,    A__					:   List[Any]					):
  '''simple docstring'''





  lowerCAmelCase_			:      Optional[int]			=   Accelerator(dispatch_batches=A__							,    split_batches=A__					)
  lowerCAmelCase_			:      Dict			=   get_dataloader(A__							,    not dispatch_batches					)
  lowerCAmelCase_			:      Dict			=   AutoModelForSequenceClassification.from_pretrained(
      """hf-internal-testing/mrpc-bert-base-cased"""							,    return_dict=A__					)
  lowerCAmelCase_,							lowerCAmelCase_			:      Dict			=   accelerator.prepare(A__							,    A__					)
  return {"ddp": [ddp_model, ddp_dataloader, "cuda:0"], "no": [model, dataloader, accelerator.device]}, accelerator







def 				UpperCamelCase_			(    A__					:   Optional[Any]							,    A__					:   int							,    A__					:   Any					):
  '''simple docstring'''





  lowerCAmelCase_			:      Optional[Any]			=   []
  for batch in dataloader:
    lowerCAmelCase_,							lowerCAmelCase_			:      Union[str, Any]			=   batch.values()
    with torch.no_grad():
      lowerCAmelCase_			:      Optional[Any]			=   model(A__					)
      lowerCAmelCase_,							lowerCAmelCase_			:      Union[str, Any]			=   accelerator.gather_for_metrics((logit, target)					)
      logits_and_targets.append((logit, target)					)
  lowerCAmelCase_,							lowerCAmelCase_			:      Optional[Any]			=   [], []
  for logit, targ in logits_and_targets:
    logits.append(A__					)
    targs.append(A__					)
  lowerCAmelCase_,							lowerCAmelCase_			:      Optional[Any]			=   torch.cat(A__					), torch.cat(A__					)
  return logits, targs







def 				UpperCamelCase_			(    A__					:   Accelerator							,    A__					:   List[str]=82							,    A__					:   int=False							,    A__					:   Tuple=False							,    A__					:   Tuple=16					):
  '''simple docstring'''





  lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_			:      Tuple			=   get_basic_setup(A__							,    A__							,    A__					)
  lowerCAmelCase_,							lowerCAmelCase_			:      Tuple			=   generate_predictions(A__							,    A__							,    A__					)
  assert (
      len(A__					) == num_samples
  ), f'Unexpected number of inputs:\n    Expected: {num_samples}\n    Actual: {len(A__					)}'







def 				UpperCamelCase_			(    A__					:   bool = False							,    A__					:   bool = False					):
  '''simple docstring'''





  lowerCAmelCase_			:      Any			=   evaluate.load("""glue"""							,    """mrpc"""					)
  lowerCAmelCase_,							lowerCAmelCase_			:      List[str]			=   get_mrpc_setup(A__							,    A__					)
  # First do baseline
  lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_			:      List[str]			=   setup["""no"""]
  model.to(A__					)
  model.eval()
  for batch in dataloader:
    batch.to(A__					)
    with torch.inference_mode():
      lowerCAmelCase_			:      Dict			=   model(**A__					)
    lowerCAmelCase_			:      Tuple			=   outputs.logits.argmax(dim=-1					)
    metric.add_batch(predictions=A__							,    references=batch["""labels"""]					)
  lowerCAmelCase_			:      Optional[Any]			=   metric.compute()

  # Then do distributed
  lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_			:      Dict			=   setup["""ddp"""]
  model.eval()
  for batch in dataloader:
    with torch.inference_mode():
      lowerCAmelCase_			:      Tuple			=   model(**A__					)
    lowerCAmelCase_			:      Union[str, Any]			=   outputs.logits.argmax(dim=-1					)
    lowerCAmelCase_			:      Any			=   batch["""labels"""]
    lowerCAmelCase_,							lowerCAmelCase_			:      Optional[Any]			=   accelerator.gather_for_metrics((preds, references)					)
    metric.add_batch(predictions=A__							,    references=A__					)
  lowerCAmelCase_			:      Any			=   metric.compute()

  for key in "accuracy f1".split():
    assert math.isclose(
        baseline[key]							,    distributed[key]					), f'Baseline and Distributed are not the same for key {key}:\n\tBaseline: {baseline[key]}\n\tDistributed: {distributed[key]}\n'







def 				UpperCamelCase_			(    ):
  '''simple docstring'''





  lowerCAmelCase_			:      Tuple			=   Accelerator(split_batches=A__							,    dispatch_batches=A__					)
  if accelerator.is_local_main_process:
    datasets.utils.logging.set_verbosity_warning()
    transformers.utils.logging.set_verbosity_warning()
  else:
    datasets.utils.logging.set_verbosity_error()
    transformers.utils.logging.set_verbosity_error()
  # These are a bit slower so they should only be ran on the GPU or TPU
  if torch.cuda.is_available() or is_tpu_available():
    if accelerator.is_local_main_process:
      print("""**Testing gather_for_metrics**"""					)
    for split_batches in [True, False]:
      for dispatch_batches in [True, False]:
        if accelerator.is_local_main_process:
          print(f'With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`'					)
        test_mrpc(A__							,    A__					)
        accelerator.state._reset_state()
  if accelerator.is_local_main_process:
    print("""**Test torch metrics**"""					)
  for split_batches in [True, False]:
    for dispatch_batches in [True, False]:
      lowerCAmelCase_			:      Tuple			=   Accelerator(split_batches=A__							,    dispatch_batches=A__					)
      if accelerator.is_local_main_process:
        print(f'With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`, length=99'					)
      test_torch_metrics(A__							,    99					)
      accelerator.state._reset_state()
  if accelerator.is_local_main_process:
    print("""**Test last batch is not dropped when perfectly divisible**"""					)
  lowerCAmelCase_			:      Any			=   Accelerator()
  test_torch_metrics(A__							,    5_12					)
  accelerator.state._reset_state()







def 				UpperCamelCase_			(    A__					:   List[Any]					):
  '''simple docstring'''





  main()


if __name__ == "__main__":
  main()







'''simple docstring'''






import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoFeatureExtractor, WavaVecaFeatureExtractor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / "utils"))

from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


__A					:  str											=      get_tests_dir("fixtures")






class    __snake_case       (   unittest.TestCase):







    """simple docstring"""



    def __lowercase   (   self    :  int		)   ->  List[str]:
      # A mock response for an HTTP head request to emulate server down
      lowerCAmelCase_			:      str			=   mock.Mock()
      lowerCAmelCase_			:      Optional[Any]			=   5_00
      lowerCAmelCase_			:      List[str]			=   {}
      lowerCAmelCase_			:      str			=   HTTPError
      lowerCAmelCase_			:      Any			=   {}

      # Download this model to make sure it's in the cache.
      lowerCAmelCase_			:      Optional[Any]			=   WavaVecaFeatureExtractor.from_pretrained("""hf-internal-testing/tiny-random-wav2vec2"""		)
      # Under the mock environment we get a 500 error when trying to reach the model.
      with mock.patch("""requests.Session.request"""    ,						return_value=lowerCamelCase		) as mock_head:
        lowerCAmelCase_			:      Union[str, Any]			=   WavaVecaFeatureExtractor.from_pretrained("""hf-internal-testing/tiny-random-wav2vec2"""		)
        # This check we did call the fake head request
        mock_head.assert_called()





    def __lowercase   (   self    :  Dict		)   ->  Any:
      # This test is for deprecated behavior and can be removed in v5
      lowerCAmelCase_			:      List[str]			=   WavaVecaFeatureExtractor.from_pretrained(
          """https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json"""		)







@is_staging_test
class    __snake_case       (   unittest.TestCase):







    """simple docstring"""



    @classmethod
    def __lowercase   (   cls    :  Tuple		)   ->  str:
      lowerCAmelCase_			:      Dict			=   TOKEN
      HfFolder.save_token(lowerCamelCase		)

    @classmethod
    def __lowercase   (   cls    :  Any		)   ->  Any:
      try:
        delete_repo(token=cls._token    ,						repo_id="""test-feature-extractor"""		)
      except HTTPError:
        pass

      try:
        delete_repo(token=cls._token    ,						repo_id="""valid_org/test-feature-extractor-org"""		)
      except HTTPError:
        pass

      try:
        delete_repo(token=cls._token    ,						repo_id="""test-dynamic-feature-extractor"""		)
      except HTTPError:
        pass

    def __lowercase   (   self    :  int		)   ->  str:
      lowerCAmelCase_			:      Tuple			=   WavaVecaFeatureExtractor.from_pretrained(lowerCamelCase		)
      feature_extractor.push_to_hub("""test-feature-extractor"""    ,						use_auth_token=self._token		)

      lowerCAmelCase_			:      Optional[Any]			=   WavaVecaFeatureExtractor.from_pretrained(F'{USER}/test-feature-extractor'		)
      for k, v in feature_extractor.__dict__.items():
        self.assertEqual(lowerCamelCase    ,						getattr(lowerCamelCase    ,						lowerCamelCase		)		)

      # Reset repo
      delete_repo(token=self._token    ,						repo_id="""test-feature-extractor"""		)

      # Push to hub via save_pretrained
      with tempfile.TemporaryDirectory() as tmp_dir:
        feature_extractor.save_pretrained(
            lowerCamelCase    ,						repo_id="""test-feature-extractor"""    ,						push_to_hub=lowerCamelCase    ,						use_auth_token=self._token		)

      lowerCAmelCase_			:      Any			=   WavaVecaFeatureExtractor.from_pretrained(F'{USER}/test-feature-extractor'		)
      for k, v in feature_extractor.__dict__.items():
        self.assertEqual(lowerCamelCase    ,						getattr(lowerCamelCase    ,						lowerCamelCase		)		)

    def __lowercase   (   self    :  Optional[Any]		)   ->  int:
      lowerCAmelCase_			:      Union[str, Any]			=   WavaVecaFeatureExtractor.from_pretrained(lowerCamelCase		)
      feature_extractor.push_to_hub("""valid_org/test-feature-extractor"""    ,						use_auth_token=self._token		)

      lowerCAmelCase_			:      Optional[int]			=   WavaVecaFeatureExtractor.from_pretrained("""valid_org/test-feature-extractor"""		)
      for k, v in feature_extractor.__dict__.items():
        self.assertEqual(lowerCamelCase    ,						getattr(lowerCamelCase    ,						lowerCamelCase		)		)

      # Reset repo
      delete_repo(token=self._token    ,						repo_id="""valid_org/test-feature-extractor"""		)

      # Push to hub via save_pretrained
      with tempfile.TemporaryDirectory() as tmp_dir:
        feature_extractor.save_pretrained(
            lowerCamelCase    ,						repo_id="""valid_org/test-feature-extractor-org"""    ,						push_to_hub=lowerCamelCase    ,						use_auth_token=self._token		)

      lowerCAmelCase_			:      Optional[int]			=   WavaVecaFeatureExtractor.from_pretrained("""valid_org/test-feature-extractor-org"""		)
      for k, v in feature_extractor.__dict__.items():
        self.assertEqual(lowerCamelCase    ,						getattr(lowerCamelCase    ,						lowerCamelCase		)		)





    def __lowercase   (   self    :  Optional[Any]		)   ->  Any:
      CustomFeatureExtractor.register_for_auto_class()
      lowerCAmelCase_			:      Dict			=   CustomFeatureExtractor.from_pretrained(lowerCamelCase		)

      feature_extractor.push_to_hub("""test-dynamic-feature-extractor"""    ,						use_auth_token=self._token		)

      # This has added the proper auto_map field to the config
      self.assertDictEqual(
          feature_extractor.auto_map    ,						{"""AutoFeatureExtractor""": """custom_feature_extraction.CustomFeatureExtractor"""}    ,						)

      lowerCAmelCase_			:      Union[str, Any]			=   AutoFeatureExtractor.from_pretrained(
          F'{USER}/test-dynamic-feature-extractor'    ,						trust_remote_code=lowerCamelCase		)
      # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
      self.assertEqual(new_feature_extractor.__class__.__name__    ,						"""CustomFeatureExtractor"""		)



