3c3
< # Copyright 2021 The HuggingFace Team. All rights reserved.
---
> # Copyright 2022 The HuggingFace Team. All rights reserved.
30a31
> import transformers
33,34d33
< 
< import transformers
45,47c44
<     Seq2SeqTrainer,
<     Seq2SeqTrainingArguments,
<     set_seed,
---
>     default_data_collator,
52a50,52
> from optimum.habana import GaudiConfig, GaudiSeq2SeqTrainer, GaudiSeq2SeqTrainingArguments
> from optimum.habana.utils import set_seed
> 
55c55
< check_min_version("4.29.0.dev0")
---
> check_min_version("4.28.0")
119a120,128
>     use_cache: bool = field(
>         default=True,
>         metadata={
>             "help": (
>                 "Whether or not the model should return the last key/values attentions (not used by all models)."
>                 "Only relevant if `config.is_decoder=True`."
>             )
>         },
>     )
203c212
<                 "efficient on GPU but very bad for TPU."
---
>                 "efficient on GPU but very bad for HPU in lazy mode."
307c316
<     parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
---
>     parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiSeq2SeqTrainingArguments))
336a346,352
>     gaudi_config = GaudiConfig.from_pretrained(
>         training_args.gaudi_config_name,
>         cache_dir=model_args.cache_dir,
>         revision=model_args.model_revision,
>         use_auth_token=True if model_args.use_auth_token else None,
>     )
> 
339,340c355,357
<         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}"
---
>         f"Process rank: {training_args.local_rank}, device: {training_args.device}, "
>         + f"distributed training: {bool(training_args.local_rank != -1)}, "
>         + f"mixed-precision training: {gaudi_config.use_habana_mixed_precision}"
420a438
>         use_cache=False if training_args.gradient_checkpointing else model_args.use_cache,
611,616c629,637
<     data_collator = DataCollatorForSeq2Seq(
<         tokenizer,
<         model=model,
<         label_pad_token_id=label_pad_token_id,
<         pad_to_multiple_of=8 if training_args.fp16 else None,
<     )
---
>     if data_args.pad_to_max_length:
>         data_collator = default_data_collator
>     else:
>         data_collator = DataCollatorForSeq2Seq(
>             tokenizer,
>             model=model,
>             label_pad_token_id=label_pad_token_id,
>             pad_to_multiple_of=8 if training_args.fp16 else None,
>         )
661c682
<     trainer = Seq2SeqTrainer(
---
>     trainer = GaudiSeq2SeqTrainer(
662a684
>         gaudi_config=gaudi_config,
745,749d766
< 
< 
< def _mp_fn(index):
<     # For xla_spawn (TPUs)
<     main()
