Update README.md

README.md

@@ -78,8 +78,8 @@ def create_optimizer(opt_model, lr_ratio=0.1):
             "lr": training_args.learning_rate * lr_ratio
         },
     ]
-    - optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)
-    + optimizer_cls, optimizer_kwargs = GaudiTrainer.get_optimizer_cls_and_kwargs(training_args)
+-   optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)
++   optimizer_cls, optimizer_kwargs = GaudiTrainer.get_optimizer_cls_and_kwargs(training_args)
     optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
 
     return optimizer
@@ -104,8 +104,8 @@ def preprocess_logits_for_metrics(logits, labels):
 
 
 if __name__ == "__main__":
-    - device = torch.device("cpu")
-    + device = torch.device("hpu")
+-   device = torch.device("cpu")
++   device = torch.device("hpu")
     raw_dataset = load_dataset("Jiqing/ProtST-BinaryLocalization")
     model = AutoModel.from_pretrained("Jiqing/protst-esm1b-for-sequential-classification", trust_remote_code=True, torch_dtype=torch.bfloat16).to(device)
     tokenizer = AutoTokenizer.from_pretrained("facebook/esm1b_t33_650M_UR50S")
@@ -117,8 +117,8 @@ if __name__ == "__main__":
                      'dataloader_num_workers': 0, 'run_name': 'downstream_esm1b_localization_fix', 'optim': 'adamw_torch', 'resume_from_checkpoint': False, \
                      - 'label_names': ['labels'], 'load_best_model_at_end': True, 'metric_for_best_model': 'accuracy', 'bf16': True, "save_total_limit": 3}
                      + 'label_names': ['labels'], 'load_best_model_at_end': True, 'metric_for_best_model': 'accuracy', 'bf16': True, "save_total_limit": 3, "use_habana":True, "use_lazy_mode": True, "use_hpu_graphs_for_inference": True}
-    - training_args = HfArgumentParser(TrainingArguments).parse_dict(training_args, allow_extra_keys=False)[0]
-    + training_args = HfArgumentParser(GaudiTrainingArguments).parse_dict(training_args, allow_extra_keys=False)[0]
+-   training_args = HfArgumentParser(TrainingArguments).parse_dict(training_args, allow_extra_keys=False)[0]
++   training_args = HfArgumentParser(GaudiTrainingArguments).parse_dict(training_args, allow_extra_keys=False)[0]
 
     def tokenize_protein(example, tokenizer=None):
         protein_seq = example["prot_seq"]
@@ -134,8 +134,8 @@ if __name__ == "__main__":
     for split in ["train", "validation", "test"]:
         raw_dataset[split] = raw_dataset[split].map(func_tokenize_protein, batched=False, remove_columns=["Unnamed: 0", "prot_seq", "localization"])
 
-    - data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
-    + data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="max_length", max_length=1024)
+-   data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
++   data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="max_length", max_length=1024)
 
     transformers.utils.logging.set_verbosity_info()
     log_level = training_args.get_process_log_level()
@@ -144,16 +144,16 @@ if __name__ == "__main__":
     optimizer = create_optimizer(model)
     scheduler = create_scheduler(training_args, optimizer)
 
-    + gaudi_config = GaudiConfig()
-    + gaudi_config.use_fused_adam = True
-    + gaudi_config.use_fused_clip_norm =True
++   gaudi_config = GaudiConfig()
++   gaudi_config.use_fused_adam = True
++   gaudi_config.use_fused_clip_norm =True
 
 
     # build trainer
-    - trainer = Trainer(
-    + trainer = GaudiTrainer(
+-   trainer = Trainer(
++   trainer = GaudiTrainer(
         model=model,
-        + gaudi_config=gaudi_config,
++       gaudi_config=gaudi_config,
         args=training_args,
         train_dataset=raw_dataset["train"],
         eval_dataset=raw_dataset["validation"],