Update README.md

README.md

@@ -5,7 +5,9 @@ tags: []
 
 # Model Card for Model ID
 
-ProtST for binary localization
+ProtST for binary localization.
+
+The following script shows how to finetune ProtST on Gaudi.
 
 ## Running script
 ```python
@@ -22,6 +24,9 @@ import torch
 import logging
 import datasets
 import transformers
++ import habana_frameworks.torch
++ from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments
+
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -73,7 +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 = 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
@@ -98,7 +104,8 @@ def preprocess_logits_for_metrics(logits, labels):
 
 
 if __name__ == "__main__":
-    device = torch.device("cpu")
+    - 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")
@@ -108,8 +115,10 @@ if __name__ == "__main__":
                      'learning_rate': 5e-05, 'weight_decay': 0, 'num_train_epochs': 100, 'max_steps': -1, 'lr_scheduler_type': 'constant', 'do_eval': True, \
                      'evaluation_strategy': 'epoch', 'per_device_eval_batch_size': 32, 'logging_strategy': 'epoch', 'save_strategy': 'epoch', 'save_steps': 820, \
                      '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}
-    training_args = HfArgumentParser(TrainingArguments).parse_dict(training_args, allow_extra_keys=False)[0]
+                     - '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]
 
     def tokenize_protein(example, tokenizer=None):
         protein_seq = example["prot_seq"]
@@ -125,7 +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)
+    + 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()
@@ -134,9 +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
+
+
     # build trainer
-    trainer = Trainer(
+    - trainer = Trainer(
+    + trainer = GaudiTrainer(
         model=model,
+        + gaudi_config=gaudi_config,
         args=training_args,
         train_dataset=raw_dataset["train"],
         eval_dataset=raw_dataset["validation"],