25d24
< import warnings
33,34d31
< from datasets import DatasetDict, load_dataset
< 
35a33
> from datasets import DatasetDict, load_dataset
43,44d40
<     Trainer,
<     TrainingArguments,
48c44
< from transformers.trainer_utils import get_last_checkpoint, is_main_process
---
> from transformers.trainer_utils import get_last_checkpoint
51a48,50
> from optimum.graphcore import IPUConfig, IPUTrainer
> from optimum.graphcore import IPUTrainingArguments as TrainingArguments
> 
104c103
<         default=0.05,
---
>         default=0.0,
292c291
<     processor: AutoProcessor
---
>     processor: Union[AutoProcessor, Wav2Vec2Processor]
324c323
<         return batch
---
>         return batch.data
349c348,349
<         lambda vocab_1, vocab_2: set(vocab_1["vocab"][0]) | set(vocab_2["vocab"][0]), vocabs.values()
---
>         lambda vocab_1, vocab_2: set(vocab_1["vocab"][0]) | set(vocab_2["vocab"][0]),
>         vocabs.values(),
381a382,384
>     if training_args.gradient_checkpointing:
>         raise ValueError("Gradient checkpointing not supported.")
> 
407d409
<     logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)
409,416c411,417
<     # Log on each process the small summary:
<     logger.warning(
<         f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
<         f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
<     )
<     # Set the verbosity to info of the Transformers logger (on main process only):
<     if is_main_process(training_args.local_rank):
<         transformers.utils.logging.set_verbosity_info()
---
>     log_level = training_args.get_process_log_level()
>     logger.setLevel(log_level)
>     datasets.utils.logging.set_verbosity(log_level)
>     transformers.utils.logging.set_verbosity(log_level)
>     transformers.utils.logging.enable_default_handler()
>     transformers.utils.logging.enable_explicit_format()
> 
493c494,496
<         model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_auth_token=data_args.use_auth_token
---
>         model_args.model_name_or_path,
>         cache_dir=model_args.cache_dir,
>         use_auth_token=data_args.use_auth_token,
551a555
> 
553c557,565
<         model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_auth_token=data_args.use_auth_token
---
>         model_args.model_name_or_path,
>         cache_dir=model_args.cache_dir,
>         use_auth_token=data_args.use_auth_token,
>     )
> 
>     ipu_config = IPUConfig.from_pretrained(
>         training_args.ipu_config_name if training_args.ipu_config_name else model_args.model_name_or_path,
>         cache_dir=model_args.cache_dir,
>         use_auth_token=True if data_args.use_auth_token else None,
572a585,586
>             "layer_norm_eps": 0.0001,
>             "apply_spec_augment": False,  # spec_augment not currently supported
585,586c599,600
<     if model_args.freeze_feature_encoder:
<         model.freeze_feature_encoder()
---
>     if not model_args.freeze_feature_encoder:
>         raise NotImplementedError("IPU version of this model freezes the feature encoder. Must be set to True.")
597c611,612
<             data_args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
---
>             data_args.audio_column_name,
>             datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate),
611c626,627
<     def prepare_dataset(batch):
---
> 
>     def prepare_dataset(batch, feature_extractor, tokenizer):
618a635,638
>         if not training_args.fp32:
>             # Cast audio inputs to FP16
>             batch["input_values"] = batch["input_values"].astype(np.float16)
> 
624a645
> 
629c650
<             prepare_dataset,
---
>             lambda batch: prepare_dataset(batch, feature_extractor, tokenizer),
675,695c696,701
<     # Now save everything to be able to create a single processor later
<     # make sure all processes wait until data is saved
<     with training_args.main_process_first():
<         # only the main process saves them
<         if is_main_process(training_args.local_rank):
<             # save feature extractor, tokenizer and config
<             feature_extractor.save_pretrained(training_args.output_dir)
<             tokenizer.save_pretrained(training_args.output_dir)
<             config.save_pretrained(training_args.output_dir)
< 
<     try:
<         processor = AutoProcessor.from_pretrained(training_args.output_dir)
<     except (OSError, KeyError):
<         warnings.warn(
<             "Loading a processor from a feature extractor config that does not"
<             " include a `processor_class` attribute is deprecated and will be removed in v5. Please add the following "
<             " attribute to your `preprocessor_config.json` file to suppress this warning: "
<             " `'processor_class': 'Wav2Vec2Processor'`",
<             FutureWarning,
<         )
<         processor = Wav2Vec2Processor.from_pretrained(training_args.output_dir)
---
>     # save feature extractor, tokenizer and config
>     feature_extractor.save_pretrained(training_args.output_dir)
>     tokenizer.save_pretrained(training_args.output_dir)
>     config.save_pretrained(training_args.output_dir)
> 
>     processor = Wav2Vec2Processor(feature_extractor, tokenizer)
698c704,708
<     data_collator = DataCollatorCTCWithPadding(processor=processor)
---
>     data_collator = DataCollatorCTCWithPadding(
>         processor=processor,
>         pad_to_multiple_of=int(max_input_length),
>         pad_to_multiple_of_labels=500,
>     )
701c711
<     trainer = Trainer(
---
>     trainer = IPUTrainer(
702a713
>         ipu_config=ipu_config,
708d718
<         tokenizer=feature_extractor,
714a725
>         logger.info("*** Train ***")