28,29d27
< from datasets import DatasetDict, load_dataset
< 
31,39c29,30
< from transformers import (
<     AutoConfig,
<     AutoFeatureExtractor,
<     AutoModelForAudioClassification,
<     HfArgumentParser,
<     Trainer,
<     TrainingArguments,
<     set_seed,
< )
---
> from datasets import DatasetDict, load_dataset
> from transformers import AutoConfig, AutoFeatureExtractor, AutoModelForAudioClassification, HfArgumentParser, set_seed
43a35,37
> from optimum.graphcore import IPUConfig, IPUTrainer
> from optimum.graphcore import IPUTrainingArguments as TrainingArguments
> 
223,226d216
<     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}"
<     )
292,316d281
<     model_input_name = feature_extractor.model_input_names[0]
< 
<     def train_transforms(batch):
<         """Apply train_transforms across a batch."""
<         subsampled_wavs = []
<         for audio in batch[data_args.audio_column_name]:
<             wav = random_subsample(
<                 audio["array"], max_length=data_args.max_length_seconds, sample_rate=feature_extractor.sampling_rate
<             )
<             subsampled_wavs.append(wav)
<         inputs = feature_extractor(subsampled_wavs, sampling_rate=feature_extractor.sampling_rate)
<         output_batch = {model_input_name: inputs.get(model_input_name)}
<         output_batch["labels"] = list(batch[data_args.label_column_name])
< 
<         return output_batch
< 
<     def val_transforms(batch):
<         """Apply val_transforms across a batch."""
<         wavs = [audio["array"] for audio in batch[data_args.audio_column_name]]
<         inputs = feature_extractor(wavs, sampling_rate=feature_extractor.sampling_rate)
<         output_batch = {model_input_name: inputs.get(model_input_name)}
<         output_batch["labels"] = list(batch[data_args.label_column_name])
< 
<         return output_batch
< 
344a310,320
> 
>     # Disable the feature that can't run on IPU yet
>     # TODO handle this properly
>     config.mask_time_prob = 0.0
> 
>     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,
>         revision=model_args.model_revision,
>         use_auth_token=True if model_args.use_auth_token else None,
>     )
358a335,353
>     def preprocess_function(examples):
>         audio = examples[data_args.audio_column_name]
>         subsampled_audio = random_subsample(
>             audio["array"], max_length=data_args.max_length_seconds, sample_rate=feature_extractor.sampling_rate
>         )
>         max_length = int(round(feature_extractor.sampling_rate * data_args.max_length_seconds))
>         inputs = feature_extractor(
>             subsampled_audio,
>             max_length=max_length,
>             sampling_rate=feature_extractor.sampling_rate,
>             padding="max_length",
>         )
>         examples["input_values"] = inputs["input_values"][0]
>         if not training_args.fp32:
>             # Cast audio input to FP16
>             examples["input_values"] = examples["input_values"].astype(np.float16)
>         examples["labels"] = examples[data_args.label_column_name]
>         return examples
> 
365c360,362
<         raw_datasets["train"].set_transform(train_transforms, output_all_columns=False)
---
>         raw_datasets["train"] = raw_datasets["train"].map(
>             preprocess_function, remove_columns=raw_datasets["train"].column_names
>         )
373c370,372
<         raw_datasets["eval"].set_transform(val_transforms, output_all_columns=False)
---
>         raw_datasets["eval"] = raw_datasets["eval"].map(
>             preprocess_function, remove_columns=raw_datasets["eval"].column_names
>         )
376c375
<     trainer = Trainer(
---
>     trainer = IPUTrainer(
377a377
>         ipu_config=ipu_config,