34,35d33
< from datasets import load_dataset
< 
36a35
> from datasets import load_dataset
43d41
<     DataCollatorForLanguageModeling,
45,47d42
<     Trainer,
<     TrainingArguments,
<     is_torch_tpu_available,
53a49,53
> from optimum.graphcore import IPUConfig, IPUTrainer
> from optimum.graphcore import IPUTrainingArguments as TrainingArguments
> from optimum.graphcore.data import DataCollatorForLanguageModelingWithMaxTokensMasked
> from optimum.graphcore.utils import check_min_version as gc_check_min_version
> 
57a58,60
> # Will error if the minimal version of Optimum Graphcore is not installed. Remove at your own risks.
> gc_check_min_version("0.6.0.dev0")
> 
209d211
<     streaming: bool = field(default=False, metadata={"help": "Enable streaming mode"})
212,214d213
<         if self.streaming:
<             require_version("datasets>=2.0.0", "The streaming feature requires `datasets>=2.0.0`")
< 
263,268d261
<     # 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):
305d297
<             streaming=data_args.streaming,
314d305
<                 streaming=data_args.streaming,
322d312
<                 streaming=data_args.streaming,
382a373,379
>     if training_args.ipu_config_name:
>         ipu_config = IPUConfig.from_pretrained(training_args.ipu_config_name, **config_kwargs)
>     elif model_args.model_name_or_path:
>         ipu_config = IPUConfig.from_pretrained(model_args.model_name_or_path, **config_kwargs)
>     else:
>         raise RuntimeError("You must provide an IPUConfig")
> 
464,478c461,468
<             if not data_args.streaming:
<                 tokenized_datasets = raw_datasets.map(
<                     tokenize_function,
<                     batched=True,
<                     num_proc=data_args.preprocessing_num_workers,
<                     remove_columns=[text_column_name],
<                     load_from_cache_file=not data_args.overwrite_cache,
<                     desc="Running tokenizer on dataset line_by_line",
<                 )
<             else:
<                 tokenized_datasets = raw_datasets.map(
<                     tokenize_function,
<                     batched=True,
<                     remove_columns=[text_column_name],
<                 )
---
>             tokenized_datasets = raw_datasets.map(
>                 tokenize_function,
>                 batched=True,
>                 num_proc=data_args.preprocessing_num_workers,
>                 remove_columns=[text_column_name],
>                 load_from_cache_file=not data_args.overwrite_cache,
>                 desc="Running tokenizer on dataset line_by_line",
>             )
487,501c477,484
<             if not data_args.streaming:
<                 tokenized_datasets = raw_datasets.map(
<                     tokenize_function,
<                     batched=True,
<                     num_proc=data_args.preprocessing_num_workers,
<                     remove_columns=column_names,
<                     load_from_cache_file=not data_args.overwrite_cache,
<                     desc="Running tokenizer on every text in dataset",
<                 )
<             else:
<                 tokenized_datasets = raw_datasets.map(
<                     tokenize_function,
<                     batched=True,
<                     remove_columns=column_names,
<                 )
---
>             tokenized_datasets = raw_datasets.map(
>                 tokenize_function,
>                 batched=True,
>                 num_proc=data_args.preprocessing_num_workers,
>                 remove_columns=column_names,
>                 load_from_cache_file=not data_args.overwrite_cache,
>                 desc="Running tokenizer on every text in dataset",
>             )
528,540c511,517
<             if not data_args.streaming:
<                 tokenized_datasets = tokenized_datasets.map(
<                     group_texts,
<                     batched=True,
<                     num_proc=data_args.preprocessing_num_workers,
<                     load_from_cache_file=not data_args.overwrite_cache,
<                     desc=f"Grouping texts in chunks of {max_seq_length}",
<                 )
<             else:
<                 tokenized_datasets = tokenized_datasets.map(
<                     group_texts,
<                     batched=True,
<                 )
---
>             tokenized_datasets = tokenized_datasets.map(
>                 group_texts,
>                 batched=True,
>                 num_proc=data_args.preprocessing_num_workers,
>                 load_from_cache_file=not data_args.overwrite_cache,
>                 desc=f"Grouping texts in chunks of {max_seq_length}",
>             )
580,581c557,558
<     pad_to_multiple_of_8 = data_args.line_by_line and training_args.fp16 and not data_args.pad_to_max_length
<     data_collator = DataCollatorForLanguageModeling(
---
>     data_collator = DataCollatorForLanguageModelingWithMaxTokensMasked(
>         max_seq_length,
584c561
<         pad_to_multiple_of=8 if pad_to_multiple_of_8 else None,
---
>         pad_to_multiple_of=None,
585a563,569
>     model.config.max_num_masked_tokens = data_collator.max_num_masked_tokens
> 
>     if training_args.do_eval and not training_args.prediction_loss_only:
>         logging.warning(
>             "Because pipelined models return only the loss sometimes (due to performance reasons), evaluation might not"
>             " work as expected, set --prediction_loss_only to fix that."
>         )
588c572
<     trainer = Trainer(
---
>     trainer = IPUTrainer(
589a574
>         ipu_config=ipu_config,
595,598c580,581
<         compute_metrics=compute_metrics if training_args.do_eval and not is_torch_tpu_available() else None,
<         preprocess_logits_for_metrics=preprocess_logits_for_metrics
<         if training_args.do_eval and not is_torch_tpu_available()
<         else None,
---
>         compute_metrics=compute_metrics if training_args.do_eval else None,
>         preprocess_logits_for_metrics=preprocess_logits_for_metrics if training_args.do_eval else None,