Update README.md

README.md

@@ -5,10 +5,12 @@ 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
+```diff
 from transformers import AutoModel, AutoTokenizer, HfArgumentParser, TrainingArguments, Trainer
 from transformers.data.data_collator import DataCollatorWithPadding
 from transformers.trainer_pt_utils import get_parameter_names
@@ -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,9 +104,10 @@ def preprocess_logits_for_metrics(logits, labels):
 
 
 if __name__ == "__main__":
-    device = torch.device("cpu")
-    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)
+-   device = torch.device("cpu")
++   device = torch.device("hpu")
+    raw_dataset = load_dataset("mila-intel/ProtST-BinaryLocalization")
+    model = AutoModel.from_pretrained("mila-intel/protst-esm1b-for-sequential-classification", trust_remote_code=True, torch_dtype=torch.bfloat16).to(device)
     tokenizer = AutoTokenizer.from_pretrained("facebook/esm1b_t33_650M_UR50S")
 
     output_dir = "/home/jiqingfe/protst/protst_2/ProtST-HuggingFace/output_dir/ProtSTModel/default/ESM-1b_PubMedBERT-abs/240123_015856"
@@ -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"],
@@ -165,6 +182,3 @@ if __name__ == "__main__":
     metric = trainer.evaluate(raw_dataset["validation"], metric_key_prefix="valid")
     print("valid metric: ", metric)
 ```
-
-
-