Update README.md

README.md

@@ -10,7 +10,7 @@ ProtST for binary localization
 ## Running script
 ```python
 from transformers import AutoModel, AutoTokenizer, HfArgumentParser, TrainingArguments, Trainer
-from transformers.data.data_collator import DataCollatorForLanguageModeling, DataCollatorForTokenClassification, DataCollatorWithPadding
+from transformers.data.data_collator import DataCollatorWithPadding
 from transformers.trainer_pt_utils import get_parameter_names
 from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 from datasets import load_dataset
@@ -125,8 +125,6 @@ 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 = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.0)
-    data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
     transformers.utils.logging.set_verbosity_info()