

import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate
# and perform gradient accumulation
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


_lowerCamelCase      :					Optional[Any]					    =      16
_lowerCamelCase      :					Union[str, Any]					    =      32





def      A__					(			__A  :						Accelerator     ,				__A  :						int = 16			)					->int:
      __A    =AutoTokenizer.from_pretrained('''bert-base-cased'''			)
      __A    =load_dataset('''glue'''     ,				'''mrpc'''			)

      def tokenize_function(__A  :						Union[str, Any]			):
            # max_length=None => use the model max length (it's actually the default)
            __A    =tokenizer(examples['''sentence1''']     ,				examples['''sentence2''']     ,				truncation=__A     ,				max_length=__A			)
            return outputs

      # Apply the method we just defined to all the examples in all the splits of the dataset
      # starting with the main process first:
      with accelerator.main_process_first():
            __A    =datasets.map(
                __A     ,				batched=__A     ,				remove_columns=['''idx''', '''sentence1''', '''sentence2''']     ,				)

      # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
      # transformers library
      __A    =tokenized_datasets.rename_column('''label'''     ,				'''labels'''			)

      def collate_fn(__A  :						Any			):
            # On TPU it's best to pad everything to the same length or training will be very slow.
            __A    =1_28 if accelerator.distributed_type == DistributedType.TPU else None
            # When using mixed precision we want round multiples of 8/16
            if accelerator.mixed_precision == "fp8":
                  __A    =16
            elif accelerator.mixed_precision != "no":
                  __A    =8
            else:
                  __A    =None

            return tokenizer.pad(
                __A     ,				padding='''longest'''     ,				max_length=__A     ,				pad_to_multiple_of=__A     ,				return_tensors='''pt'''     ,				)

      # Instantiate dataloaders.
      __A    =DataLoader(
          tokenized_datasets['''train''']     ,				shuffle=__A     ,				collate_fn=__A     ,				batch_size=__A			)
      __A    =DataLoader(
          tokenized_datasets['''validation''']     ,				shuffle=__A     ,				collate_fn=__A     ,				batch_size=__A			)

      return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('''TESTING_MOCKED_DATALOADERS''', None) == "1":
   from accelerate.test_utils.training import mocked_dataloaders

   _lowerCamelCase      :					Dict					    =      mocked_dataloaders  # noqa: F811





def      A__					(			__A  :						Any     ,				__A  :						Optional[Any]			)					->List[Any]:
      # For testing only
      if os.environ.get('''TESTING_MOCKED_DATALOADERS'''     ,				__A			) == "1":
            __A    =2
      # New Code #
      __A    =int(args.gradient_accumulation_steps			)
      # Initialize accelerator
      __A    =Accelerator(
          cpu=args.cpu     ,				mixed_precision=args.mixed_precision     ,				gradient_accumulation_steps=__A			)
      if accelerator.distributed_type == DistributedType.TPU and gradient_accumulation_steps > 1:
            raise NotImplementedError(
                '''Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`'''			)
      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      __A    =config['''lr''']
      __A    =int(config['''num_epochs''']			)
      __A    =int(config['''seed''']			)
      __A    =int(config['''batch_size''']			)

      __A    =evaluate.load('''glue'''     ,				'''mrpc'''			)

      set_seed(__A			)
      __A    ,		__A    =get_dataloaders(__A     ,				__A			)
      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      __A    =AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''     ,				return_dict=__A			)

      # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
      # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
      # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
      __A    =model.to(accelerator.device			)

      # Instantiate optimizer
      __A    =AdamW(params=model.parameters()     ,				lr=__A			)

      # Instantiate scheduler
      __A    =get_linear_schedule_with_warmup(
          optimizer=__A     ,				num_warmup_steps=1_00     ,				num_training_steps=(len(__A			) * num_epochs)     ,				)

      # Prepare everything
      # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
      # prepare method.
      __A    ,		__A    ,		__A    ,		__A    ,		__A    =accelerator.prepare(
          __A     ,				__A     ,				__A     ,				__A     ,				__A			)

      # Now we train the model
      for epoch in range(__A			):
            model.train()
            for step, batch in enumerate(__A			):
                  # We could avoid this line since we set the accelerator with `device_placement=True`.
                  batch.to(accelerator.device			)
                  # New code #
                  # We use the new `accumulate` context manager to perform gradient accumulation
                  # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
                  with accelerator.accumulate(__A			):
                        __A    =model(**__A			)
                        __A    =output.loss
                        accelerator.backward(__A			)
                        optimizer.step()
                        lr_scheduler.step()
                        optimizer.zero_grad()

            model.eval()
            for step, batch in enumerate(__A			):
                  # We could avoid this line since we set the accelerator with `device_placement=True`.
                  batch.to(accelerator.device			)
                  with torch.no_grad():
                        __A    =model(**__A			)
                  __A    =outputs.logits.argmax(dim=-1			)
                  __A    ,		__A    =accelerator.gather_for_metrics((predictions, batch['''labels'''])			)
                  metric.add_batch(
                      predictions=__A     ,				references=__A     ,				)

            __A    =metric.compute()
            # Use accelerator.print to print only on the main process.
            accelerator.print(F'''epoch {epoch}:'''     ,				__A			)





def      A__					(			)					->Optional[Any]:
      __A    =argparse.ArgumentParser(description='''Simple example of training script.'''			)
      parser.add_argument(
          '''--mixed_precision'''     ,				type=__A     ,				default=__A     ,				choices=['''no''', '''fp16''', '''bf16''', '''fp8''']     ,				help='''Whether to use mixed precision. Choose'''
          '''between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'''
          '''and an Nvidia Ampere GPU.'''     ,				)
      # New Code #
      parser.add_argument(
          '''--gradient_accumulation_steps'''     ,				type=__A     ,				default=1     ,				help='''The number of minibatches to be ran before gradients are accumulated.'''     ,				)
      parser.add_argument('''--cpu'''     ,				action='''store_true'''     ,				help='''If passed, will train on the CPU.'''			)
      __A    =parser.parse_args()
      __A    ={'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
      training_function(__A     ,				__A			)


if __name__ == "__main__":
   main()





import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pytest

from datasets.arrow_dataset import Dataset
from datasets.search import ElasticSearchIndex, FaissIndex, MissingIndex

from .utils import require_elasticsearch, require_faiss


_lowerCamelCase      :					List[str]					    =      pytest.mark.integration





@require_faiss
class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =Dataset.from_dict({'''filename''': ['''my_name-train''' + '''_''' + str(lowercase__     ) for x in np.arange(3_0     ).tolist()]}     )
             return dset

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =self._create_dummy_dataset()
             __A    =dset.map(
                 lambda lowercase__ ,   lowercase__     : {"vecs": i * np.ones(5 ,   dtype=np.floataa     )} ,   with_indices=lowercase__ ,   keep_in_memory=lowercase__     )
             __A    =dset.add_faiss_index('''vecs''' ,   batch_size=1_0_0 ,   metric_type=faiss.METRIC_INNER_PRODUCT     )
             __A    ,		__A    =dset.get_nearest_examples('''vecs''' ,   np.ones(5 ,   dtype=np.floataa     )     )
             self.assertEqual(examples['''filename'''][0] ,   '''my_name-train_29'''     )
             dset.drop_index('''vecs'''     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =self._create_dummy_dataset()
             dset.add_faiss_index_from_external_arrays(
                 external_arrays=np.ones((3_0, 5)     ) * np.arange(3_0     ).reshape(-1 ,   1     ) ,   index_name='''vecs''' ,   batch_size=1_0_0 ,   metric_type=faiss.METRIC_INNER_PRODUCT ,   )
             __A    ,		__A    =dset.get_nearest_examples('''vecs''' ,   np.ones(5 ,   dtype=np.floataa     )     )
             self.assertEqual(examples['''filename'''][0] ,   '''my_name-train_29'''     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =self._create_dummy_dataset()
             dset.add_faiss_index_from_external_arrays(
                 external_arrays=np.ones((3_0, 5)     ) * np.arange(3_0     ).reshape(-1 ,   1     ) ,   index_name='''vecs''' ,   metric_type=faiss.METRIC_INNER_PRODUCT ,   )

             # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
             # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
             # see https://bugs.python.org/issue14243 and
             # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
             with tempfile.NamedTemporaryFile(delete=lowercase__     ) as tmp_file:
                   dset.save_faiss_index('''vecs''' ,   tmp_file.name     )
                   dset.load_faiss_index('''vecs2''' ,   tmp_file.name     )
             os.unlink(tmp_file.name     )

             __A    ,		__A    =dset.get_nearest_examples('''vecs2''' ,   np.ones(5 ,   dtype=np.floataa     )     )
             self.assertEqual(examples['''filename'''][0] ,   '''my_name-train_29'''     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self._create_dummy_dataset()
             dset.add_faiss_index_from_external_arrays(
                 external_arrays=np.ones((3_0, 5)     ) * np.arange(3_0     ).reshape(-1 ,   1     ) ,   index_name='''vecs'''     )
             dset.drop_index('''vecs'''     )
             self.assertRaises(lowercase__ ,   partial(dset.get_nearest_examples ,   '''vecs2''' ,   np.ones(5 ,   dtype=np.floataa     )     )     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             from elasticsearch import Elasticsearch

             __A    =self._create_dummy_dataset()
             with patch('''elasticsearch.Elasticsearch.search'''     ) as mocked_search, patch(
                 '''elasticsearch.client.IndicesClient.create'''     ) as mocked_index_create, patch('''elasticsearch.helpers.streaming_bulk'''     ) as mocked_bulk:
                   __A    ={'''acknowledged''': True}
                   mocked_bulk.return_value([(True, None)] * 3_0     )
                   __A    ={'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 2_9}]}}
                   __A    =Elasticsearch()

                   dset.add_elasticsearch_index('''filename''' ,   es_client=lowercase__     )
                   __A    ,		__A    =dset.get_nearest_examples('''filename''' ,   '''my_name-train_29'''     )
                   self.assertEqual(examples['''filename'''][0] ,   '''my_name-train_29'''     )





@require_faiss
class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT     )

             # add vectors
             index.add_vectors(np.eye(5 ,   dtype=np.floataa     )     )
             self.assertIsNotNone(index.faiss_index     )
             self.assertEqual(index.faiss_index.ntotal ,   5     )
             index.add_vectors(np.zeros((5, 5) ,   dtype=np.floataa     )     )
             self.assertEqual(index.faiss_index.ntotal ,   1_0     )

             # single query
             __A    =np.zeros(5 ,   dtype=np.floataa     )
             __A    =1
             __A    ,		__A    =index.search(lowercase__     )
             self.assertRaises(lowercase__ ,   index.search ,   query.reshape(-1 ,   1     )     )
             self.assertGreater(scores[0] ,   0     )
             self.assertEqual(indices[0] ,   1     )

             # batched queries
             __A    =np.eye(5 ,   dtype=np.floataa     )[::-1]
             __A    ,		__A    =index.search_batch(lowercase__     )
             self.assertRaises(lowercase__ ,   index.search_batch ,   queries[0]     )
             __A    =[scores[0] for scores in total_scores]
             __A    =[indices[0] for indices in total_indices]
             self.assertGreater(np.min(lowercase__     ) ,   0     )
             self.assertListEqual([4, 3, 2, 1, 0] ,   lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =FaissIndex(string_factory='''Flat'''     )
             index.add_vectors(np.eye(5 ,   dtype=np.floataa     )     )
             self.assertIsInstance(index.faiss_index ,   faiss.IndexFlat     )
             __A    =FaissIndex(string_factory='''LSH'''     )
             index.add_vectors(np.eye(5 ,   dtype=np.floataa     )     )
             self.assertIsInstance(index.faiss_index ,   faiss.IndexLSH     )
             with self.assertRaises(lowercase__     ):
                   __A    =FaissIndex(string_factory='''Flat''' ,   custom_index=faiss.IndexFlat(5     )     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =faiss.IndexFlat(5     )
             __A    =FaissIndex(custom_index=lowercase__     )
             index.add_vectors(np.eye(5 ,   dtype=np.floataa     )     )
             self.assertIsInstance(index.faiss_index ,   faiss.IndexFlat     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT     )
             index.add_vectors(np.eye(5 ,   dtype=np.floataa     )     )

             # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
             # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
             # see https://bugs.python.org/issue14243 and
             # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
             with tempfile.NamedTemporaryFile(delete=lowercase__     ) as tmp_file:
                   index.save(tmp_file.name     )
                   __A    =FaissIndex.load(tmp_file.name     )
             os.unlink(tmp_file.name     )

             __A    =np.zeros(5 ,   dtype=np.floataa     )
             __A    =1
             __A    ,		__A    =index.search(lowercase__     )
             self.assertGreater(scores[0] ,   0     )
             self.assertEqual(indices[0] ,   1     )





@require_faiss
def      A__					(			__A  :						Dict			)					->List[Any]:
      import faiss

      __A    =FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT			)
      index.add_vectors(np.eye(5     ,				dtype=np.floataa			)			)

      __A    ='''index.faiss'''
      __A    =F'''mock://{index_name}'''
      index.save(__A     ,				storage_options=mockfs.storage_options			)
      __A    =FaissIndex.load(__A     ,				storage_options=mockfs.storage_options			)

      __A    =np.zeros(5     ,				dtype=np.floataa			)
      __A    =1
      __A    ,		__A    =index.search(__A			)
      assert scores[0] > 0
      assert indices[0] == 1





@require_elasticsearch
class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             from elasticsearch import Elasticsearch

             with patch('''elasticsearch.Elasticsearch.search'''     ) as mocked_search, patch(
                 '''elasticsearch.client.IndicesClient.create'''     ) as mocked_index_create, patch('''elasticsearch.helpers.streaming_bulk'''     ) as mocked_bulk:
                   __A    =Elasticsearch()
                   __A    ={'''acknowledged''': True}
                   __A    =ElasticSearchIndex(es_client=lowercase__     )
                   mocked_bulk.return_value([(True, None)] * 3     )
                   index.add_documents(['''foo''', '''bar''', '''foobar''']     )

                   # single query
                   __A    ='''foo'''
                   __A    ={'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 0}]}}
                   __A    ,		__A    =index.search(lowercase__     )
                   self.assertEqual(scores[0] ,   1     )
                   self.assertEqual(indices[0] ,   0     )

                   # single query with timeout
                   __A    ='''foo'''
                   __A    ={'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 0}]}}
                   __A    ,		__A    =index.search(lowercase__ ,   request_timeout=3_0     )
                   self.assertEqual(scores[0] ,   1     )
                   self.assertEqual(indices[0] ,   0     )

                   # batched queries
                   __A    =['''foo''', '''bar''', '''foobar''']
                   __A    ={'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 1}]}}
                   __A    ,		__A    =index.search_batch(lowercase__     )
                   __A    =[scores[0] for scores in total_scores]
                   __A    =[indices[0] for indices in total_indices]
                   self.assertGreater(np.min(lowercase__     ) ,   0     )
                   self.assertListEqual([1, 1, 1] ,   lowercase__     )

                   # batched queries with timeout
                   __A    =['''foo''', '''bar''', '''foobar''']
                   __A    ={'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 1}]}}
                   __A    ,		__A    =index.search_batch(lowercase__ ,   request_timeout=3_0     )
                   __A    =[scores[0] for scores in total_scores]
                   __A    =[indices[0] for indices in total_indices]
                   self.assertGreater(np.min(lowercase__     ) ,   0     )
                   self.assertListEqual([1, 1, 1] ,   lowercase__     )


