33d32
< import evaluate
35,36d33
< from datasets import load_dataset
< 
37a35
> from datasets import load_dataset
45,46d42
<     Trainer,
<     TrainingArguments,
48d43
<     is_torch_tpu_available,
55a51,53
> from optimum.graphcore import IPUConfig, IPUTrainer
> from optimum.graphcore import IPUTrainingArguments as TrainingArguments
> 
130c128
<             "choices": ["auto", "bfloat16", "float16", "float32"],
---
>             "choices": ["auto", "float16", "float32"],
143,148d140
<     def __post_init__(self):
<         if self.config_overrides is not None and (self.config_name is not None or self.model_name_or_path is not None):
<             raise ValueError(
<                 "--config_overrides can't be used in combination with --config_name or --model_name_or_path"
<             )
< 
185c177
<     streaming: bool = field(default=False, metadata={"help": "Enable streaming mode"})
---
> 
214,216d205
<         if self.streaming:
<             require_version("datasets>=2.0.0", "The streaming feature requires `datasets>=2.0.0`")
< 
263,267d251
<     # 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}"
<     )
304d287
<             streaming=data_args.streaming,
313d295
<                 streaming=data_args.streaming,
321d302
<                 streaming=data_args.streaming,
385,388c366,379
<         if model_args.config_overrides is not None:
<             logger.info(f"Overriding config: {model_args.config_overrides}")
<             config.update_from_string(model_args.config_overrides)
<             logger.info(f"New config: {config}")
---
> 
>     if model_args.config_overrides is not None:
>         logger.info(f"Overriding config: {model_args.config_overrides}")
>         config.update_from_string(model_args.config_overrides)
>         logger.info(f"New config: {config}")
> 
>     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. If the pretrained model does not contain an IPUConfig, you can load an IPUConfig using --ipu_config_name."
>         )
456,470c447,454
<         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 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 dataset",
>         )
514,526c498,504
<         if not data_args.streaming:
<             lm_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 {block_size}",
<             )
<         else:
<             lm_datasets = tokenized_datasets.map(
<                 group_texts,
<                 batched=True,
<             )
---
>         lm_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 {block_size}",
>         )
544,560d521
<         def preprocess_logits_for_metrics(logits, labels):
<             if isinstance(logits, tuple):
<                 # Depending on the model and config, logits may contain extra tensors,
<                 # like past_key_values, but logits always come first
<                 logits = logits[0]
<             return logits.argmax(dim=-1)
< 
<         metric = evaluate.load("accuracy")
< 
<         def compute_metrics(eval_preds):
<             preds, labels = eval_preds
<             # preds have the same shape as the labels, after the argmax(-1) has been calculated
<             # by preprocess_logits_for_metrics but we need to shift the labels
<             labels = labels[:, 1:].reshape(-1)
<             preds = preds[:, :-1].reshape(-1)
<             return metric.compute(predictions=preds, references=labels)
< 
562c523
<     trainer = Trainer(
---
>     trainer = IPUTrainer(
563a525
>         ipu_config=ipu_config,
568d529
<         # Data collator will default to DataCollatorWithPadding, so we change it.
570,573d530
<         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,