






import argparse
from typing import List

import evaluate
import numpy as np
import torch
from datasets import DatasetDict, load_dataset

# New Code #
# We'll be using StratifiedKFold for this example
from sklearn.model_selection import StratifiedKFold
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,
# specifically showcasing how to perform Cross Validation,
# and builds off the `nlp_example.py` script.
#
# 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 help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# 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
#
########################################################################


__snake_case												=  16
__snake_case												=  32


def   _A   (				SCREAMING_SNAKE_CASE__  :	Accelerator     ,  SCREAMING_SNAKE_CASE__  :	DatasetDict     ,  SCREAMING_SNAKE_CASE__  :	List[int]     ,  SCREAMING_SNAKE_CASE__  :	List[int]     ,  SCREAMING_SNAKE_CASE__  :	int = 16       ):
    UpperCamelCase       :List[str]								=   AutoTokenizer.from_pretrained('''bert-base-cased'''       )
    UpperCamelCase       :str								=   DatasetDict(
        {
            '''train''': dataset['''train'''].select(__UpperCAmelCase       ),
            '''validation''': dataset['''train'''].select(__UpperCAmelCase       ),
            '''test''': dataset['''validation'''],
        }       )

    def tokenize_function(SCREAMING_SNAKE_CASE__  :	Union[str, Any]       ):
        # max_length=None => use the model max length (it's actually the default)
        UpperCamelCase       :Optional[int]								=   tokenizer(examples['''sentence1''']     ,  examples['''sentence2''']     ,  truncation=__UpperCAmelCase     ,  max_length=__UpperCAmelCase       )
        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():
        UpperCamelCase       :Dict								=   datasets.map(
            __UpperCAmelCase     ,  batched=__UpperCAmelCase     ,  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
    UpperCamelCase       :str								=   tokenized_datasets.rename_column('''label'''     ,  '''labels'''       )

    def collate_fn(SCREAMING_SNAKE_CASE__  :	str       ):
        # On TPU it's best to pad everything to the same length or training will be very slow.
        UpperCamelCase       :List[str]								=   128 if accelerator.distributed_type == DistributedType.TPU else None
        # When using mixed precision we want round multiples of 8/16
        if accelerator.mixed_precision == "fp8":
            UpperCamelCase       :Union[str, Any]								=   16
        elif accelerator.mixed_precision != "no":
            UpperCamelCase       :Optional[int]								=   8
        else:
            UpperCamelCase       :str								=   None

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

    # Instantiate dataloaders.
    UpperCamelCase       :int								=   DataLoader(
        tokenized_datasets['''train''']     ,  shuffle=__UpperCAmelCase     ,  collate_fn=__UpperCAmelCase     ,  batch_size=__UpperCAmelCase       )
    UpperCamelCase       :List[str]								=   DataLoader(
        tokenized_datasets['''validation''']     ,  shuffle=__UpperCAmelCase     ,  collate_fn=__UpperCAmelCase     ,  batch_size=__UpperCAmelCase       )

    UpperCamelCase       :Union[str, Any]								=   DataLoader(
        tokenized_datasets['''test''']     ,  shuffle=__UpperCAmelCase     ,  collate_fn=__UpperCAmelCase     ,  batch_size=__UpperCAmelCase       )

    return train_dataloader, eval_dataloader, test_dataloader


def   _A   (				SCREAMING_SNAKE_CASE__  :	Union[str, Any]     ,  SCREAMING_SNAKE_CASE__  :	Any       ):
    # New Code #
    UpperCamelCase       :List[Any]								=   []
    # Download the dataset
    UpperCamelCase       :Union[str, Any]								=   load_dataset('''glue'''     ,  '''mrpc'''       )
    # Create our splits
    UpperCamelCase       :Dict								=   StratifiedKFold(n_splits=int(args.num_folds       )       )
    # Initialize accelerator
    UpperCamelCase       :Union[str, Any]								=   Accelerator(cpu=args.cpu     ,  mixed_precision=args.mixed_precision       )
    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    UpperCamelCase       :Optional[int]								=   config['''lr''']
    UpperCamelCase       :str								=   int(config['''num_epochs''']       )
    UpperCamelCase       :Tuple								=   int(config['''seed''']       )
    UpperCamelCase       :List[Any]								=   int(config['''batch_size''']       )

    UpperCamelCase       :int								=   evaluate.load('''glue'''     ,  '''mrpc'''       )

    # If the batch size is too big we use gradient accumulation
    UpperCamelCase       :Any								=   1
    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
        UpperCamelCase       :Optional[Any]								=   batch_size // MAX_GPU_BATCH_SIZE
        UpperCamelCase       :Tuple								=   MAX_GPU_BATCH_SIZE

    set_seed(__UpperCAmelCase       )

    # New Code #
    # Create our folds:
    UpperCamelCase       :Any								=   kfold.split(np.zeros(datasets['''train'''].num_rows       )     ,  datasets['''train''']['''label''']       )
    UpperCamelCase       :List[str]								=   []
    # Iterate over them
    for i, (train_idxs, valid_idxs) in enumerate(__UpperCAmelCase       ):
        UpperCamelCase							,   UpperCamelCase							,   UpperCamelCase       :Union[str, Any]								=   get_fold_dataloaders(
            __UpperCAmelCase     ,  __UpperCAmelCase     ,  __UpperCAmelCase     ,  __UpperCAmelCase     ,  )
        # Instantiate the model (we build the model here so that the seed also control new weights initialization)
        UpperCamelCase       :Tuple								=   AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''     ,  return_dict=__UpperCAmelCase       )

        # 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).
        UpperCamelCase       :str								=   model.to(accelerator.device       )

        # Instantiate optimizer
        UpperCamelCase       :Union[str, Any]								=   AdamW(params=model.parameters()     ,  lr=__UpperCAmelCase       )

        # Instantiate scheduler
        UpperCamelCase       :Tuple								=   get_linear_schedule_with_warmup(
            optimizer=__UpperCAmelCase     ,  num_warmup_steps=100     ,  num_training_steps=(len(__UpperCAmelCase       ) * num_epochs) // gradient_accumulation_steps     ,  )

        # 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.
        UpperCamelCase							,   UpperCamelCase							,   UpperCamelCase							,   UpperCamelCase							,   UpperCamelCase       :Tuple								=   accelerator.prepare(
            __UpperCAmelCase     ,  __UpperCAmelCase     ,  __UpperCAmelCase     ,  __UpperCAmelCase     ,  __UpperCAmelCase       )

        # Now we train the model
        for epoch in range(__UpperCAmelCase       ):
            model.train()
            for step, batch in enumerate(__UpperCAmelCase       ):
                # We could avoid this line since we set the accelerator with `device_placement=True`.
                batch.to(accelerator.device       )
                UpperCamelCase       :str								=   model(**__UpperCAmelCase       )
                UpperCamelCase       :Tuple								=   outputs.loss
                UpperCamelCase       :int								=   loss / gradient_accumulation_steps
                accelerator.backward(__UpperCAmelCase       )
                if step % gradient_accumulation_steps == 0:
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()

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

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

        # New Code #
        # We also run predictions on the test set at the very end
        UpperCamelCase       :Optional[int]								=   []
        for step, batch in enumerate(__UpperCAmelCase       ):
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            batch.to(accelerator.device       )
            with torch.no_grad():
                UpperCamelCase       :str								=   model(**__UpperCAmelCase       )
            UpperCamelCase       :Union[str, Any]								=   outputs.logits
            UpperCamelCase							,   UpperCamelCase       :Optional[int]								=   accelerator.gather_for_metrics((predictions, batch['''labels'''])       )
            fold_predictions.append(predictions.cpu()       )
            if i == 0:
                # We need all of the test predictions
                test_references.append(references.cpu()       )
        # Use accelerator.print to print only on the main process.
        test_predictions.append(torch.cat(__UpperCAmelCase     ,  dim=0       )       )
        # We now need to release all our memory and get rid of the current model, optimizer, etc
        accelerator.free_memory()
    # New Code #
    # Finally we check the accuracy of our folded results:
    UpperCamelCase       :Tuple								=   torch.cat(__UpperCAmelCase     ,  dim=0       )
    UpperCamelCase       :List[Any]								=   torch.stack(__UpperCAmelCase     ,  dim=0       ).sum(dim=0       ).div(int(args.num_folds       )       ).argmax(dim=-1       )
    UpperCamelCase       :Dict								=   metric.compute(predictions=__UpperCAmelCase     ,  references=__UpperCAmelCase       )
    accelerator.print('''Average test metrics from all folds:'''     ,  __UpperCAmelCase       )


def   _A   (				):
    UpperCamelCase       :List[str]								=   argparse.ArgumentParser(description='''Simple example of training script.'''       )
    parser.add_argument(
        '''--mixed_precision'''     ,  type=__UpperCAmelCase     ,  default=__UpperCAmelCase     ,  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.'''     ,  )
    parser.add_argument('''--cpu'''     ,  action='''store_true'''     ,  help='''If passed, will train on the CPU.'''       )
    # New Code #
    parser.add_argument('''--num_folds'''     ,  type=__UpperCAmelCase     ,  default=3     ,  help='''The number of splits to perform across the dataset'''       )
    UpperCamelCase       :Dict								=   parser.parse_args()
    UpperCamelCase       :Any								=   {'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
    training_function(__UpperCAmelCase     ,  __UpperCAmelCase       )


if __name__ == "__main__":
 main()











import json
import os
from typing import Dict, List, Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCamelCase__ :   List[Any]					=							logging.get_logger(__name__)


lowerCamelCase__ :   List[str]					=							{
    'vocab_file': 'vocab.json',
    'tokenizer_config_file': 'tokenizer_config.json',
    'merges_file': 'merges.txt',
}

lowerCamelCase__ :   Union[str, Any]					=							{
    'vocab_file': {
        'facebook/s2t-wav2vec2-large-en-de': (
            'https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/vocab.json'
        ),
    },
    'tokenizer_config_file': {
        'facebook/s2t-wav2vec2-large-en-de': (
            'https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/tokenizer_config.json'
        ),
    },
    'merges_file': {
        'facebook/s2t-wav2vec2-large-en-de': (
            'https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/merges.txt'
        ),
    },
}

lowerCamelCase__ :   Optional[Any]					=							'</w>'
lowerCamelCase__ :   Union[str, Any]					=							'@@ '

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	Optional[Any]							)					-> Dict:
      SCREAMING_SNAKE_CASE_							=			set()
      SCREAMING_SNAKE_CASE_							=			word[0]
      for char in word[1:]:
            pairs.add((prev_char, char)							)
            SCREAMING_SNAKE_CASE_							=			char
      return pairs


# Speech2Text2 has no max input length
lowerCamelCase__ :   Any					=							{'facebook/s2t-wav2vec2-large-en-de': 1_024}





class        lowerCamelCase_							(      _SCREAMING_SNAKE_CASE					):


 '''simple docstring'''







 lowercase_				    =    VOCAB_FILES_NAMES
 lowercase_				    =    PRETRAINED_VOCAB_FILES_MAP
 lowercase_				    =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 lowercase_				    =    ["input_ids", "attention_mask"]



 def __init__(							self  :		List[str]				,				_lowerCAmelCase  :		Union[str, Any]				,				_lowerCAmelCase  :		List[Any]="<s>"				,				_lowerCAmelCase  :		Any="<pad>"				,				_lowerCAmelCase  :		List[str]="</s>"				,				_lowerCAmelCase  :		int="<unk>"				,				_lowerCAmelCase  :		Optional[Any]=False				,				_lowerCAmelCase  :		Dict=None				,				**_lowerCAmelCase  :		Tuple				,				):
       super().__init__(
           unk_token=_lowerCAmelCase				,				bos_token=_lowerCAmelCase				,				eos_token=_lowerCAmelCase				,				pad_token=_lowerCAmelCase				,				do_lower_case=_lowerCAmelCase				,				**_lowerCAmelCase				,				)

       SCREAMING_SNAKE_CASE_							=			do_lower_case

       with open(_lowerCAmelCase				,				encoding='utf-8'						) as vocab_handle:
             SCREAMING_SNAKE_CASE_							=			json.load(_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			{v: k for k, v in self.encoder.items()}

       if merges_file is None:
             logger.info(F"No merges files provided. {self.__class__.__name__} can only be used for decoding."						)

             SCREAMING_SNAKE_CASE_							=			None
             SCREAMING_SNAKE_CASE_							=			None
       else:
             with open(_lowerCAmelCase				,				encoding='utf-8'						) as merges_handle:
                   SCREAMING_SNAKE_CASE_							=			merges_handle.read().split('\n'						)[:-1]

             SCREAMING_SNAKE_CASE_							=			[tuple(merge.split()[:2]						) for merge in merges]
             SCREAMING_SNAKE_CASE_							=			dict(zip(_lowerCAmelCase				,				range(len(_lowerCAmelCase						)						)						)						)
             SCREAMING_SNAKE_CASE_							=			{}



 @property
 def       lowerCAmelCase_  (							self  :		List[str]						):
       return len(self.decoder						)



 def       lowerCAmelCase_  (							self  :		Tuple						):
       return dict(self.encoder				,				**self.added_tokens_encoder						)



 def       lowerCAmelCase_  (							self  :		Tuple				,				_lowerCAmelCase  :		Tuple						):
       SCREAMING_SNAKE_CASE_							=			tuple(token[:-1]						) + (token[-1] + BPE_TOKEN_MERGES,)
       if token in self.cache:
             return self.cache[token]
       SCREAMING_SNAKE_CASE_							=			get_pairs(_lowerCAmelCase						)

       if not pairs:
             return token

       while True:
             SCREAMING_SNAKE_CASE_							=			min(_lowerCAmelCase				,				key=lambda _lowerCAmelCase						: self.bpe_ranks.get(_lowerCAmelCase				,				float('inf'						)						)						)
             if bigram not in self.bpe_ranks:
                   break
             SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_							=			bigram
             SCREAMING_SNAKE_CASE_							=			[]
             SCREAMING_SNAKE_CASE_							=			0
             while i < len(_lowerCAmelCase						):
                   try:
                         SCREAMING_SNAKE_CASE_							=			word.index(_lowerCAmelCase				,				_lowerCAmelCase						)
                   except ValueError:
                         new_word.extend(word[i:]						)
                         break
                   else:
                         new_word.extend(word[i:j]						)
                         SCREAMING_SNAKE_CASE_							=			j

                   if word[i] == first and i < len(_lowerCAmelCase						) - 1 and word[i + 1] == second:
                         new_word.append(first + second						)
                         i += 2
                   else:
                         new_word.append(word[i]						)
                         i += 1
             SCREAMING_SNAKE_CASE_							=			tuple(_lowerCAmelCase						)
             SCREAMING_SNAKE_CASE_							=			new_word
             if len(_lowerCAmelCase						) == 1:
                   break
             else:
                   SCREAMING_SNAKE_CASE_							=			get_pairs(_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			' '.join(_lowerCAmelCase						)
       if word == "\n  " + BPE_TOKEN_MERGES:
             SCREAMING_SNAKE_CASE_							=			'\n' + BPE_TOKEN_MERGES

       if word.endswith(_lowerCAmelCase						):
             SCREAMING_SNAKE_CASE_							=			word.replace(_lowerCAmelCase				,				''						)

       SCREAMING_SNAKE_CASE_							=			word.replace(' '				,				_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			word
       return word



 def       lowerCAmelCase_  (							self  :		Optional[int]				,				_lowerCAmelCase  :		Optional[int]						):

       if self.bpe_ranks is None:
             raise ValueError(
                 'This tokenizer was instantiated without a `merges.txt` file, so'
                 ' that it can only be used for decoding, not for encoding.'
                 'Make sure to provide `merges.txt` file at instantiation to enable '
                 'encoding.'						)

       if self.do_lower_case:
             SCREAMING_SNAKE_CASE_							=			text.lower()

       SCREAMING_SNAKE_CASE_							=			text.split()

       SCREAMING_SNAKE_CASE_							=			[]
       for token in text:
             if token:
                   split_tokens.extend(list(self.bpe(_lowerCAmelCase						).split(' '						)						)						)

       return split_tokens



 def       lowerCAmelCase_  (							self  :		Union[str, Any]				,				_lowerCAmelCase  :		str						):
       return self.encoder.get(_lowerCAmelCase				,				self.encoder.get(self.unk_token						)						)



 def       lowerCAmelCase_  (							self  :		int				,				_lowerCAmelCase  :		int						):
       SCREAMING_SNAKE_CASE_							=			self.decoder.get(_lowerCAmelCase				,				self.unk_token						)
       return result



 def       lowerCAmelCase_  (							self  :		int				,				_lowerCAmelCase  :		List[str]						):
       SCREAMING_SNAKE_CASE_							=			' '.join(_lowerCAmelCase						)

       # make sure @@ tokens are concatenated
       SCREAMING_SNAKE_CASE_							=			''.join(string.split(_lowerCAmelCase						)						)

       return string




 def       lowerCAmelCase_  (							self  :		int				,				_lowerCAmelCase  :		str				,				_lowerCAmelCase  :		Optional[str] = None						):
       if not os.path.isdir(_lowerCAmelCase						):
             logger.error(F"Vocabulary path ({save_directory}) should be a directory"						)
             return
       SCREAMING_SNAKE_CASE_							=			os.path.join(
           _lowerCAmelCase				,				(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']						)
       SCREAMING_SNAKE_CASE_							=			os.path.join(
           _lowerCAmelCase				,				(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file']						)

       with open(_lowerCAmelCase				,				'w'				,				encoding='utf-8'						) as f:
             f.write(json.dumps(self.encoder				,				indent=2				,				sort_keys=_lowerCAmelCase				,				ensure_ascii=_lowerCAmelCase						) + '\n'						)

       SCREAMING_SNAKE_CASE_							=			0
       if self.bpe_ranks is None:
             return (vocab_file,)

       with open(_lowerCAmelCase				,				'w'				,				encoding='utf-8'						) as writer:
             for bpe_tokens, token_index in sorted(self.bpe_ranks.items()				,				key=lambda _lowerCAmelCase						: kv[1]						):
                   if index != token_index:
                         logger.warning(
                             F"Saving vocabulary to {merges_file}: BPE merge indices are not consecutive."
                             ' Please check that the tokenizer is not corrupted!'						)
                         SCREAMING_SNAKE_CASE_							=			token_index
                   writer.write(' '.join(_lowerCAmelCase						) + '\n'						)
                   index += 1

       return (vocab_file, merges_file)