Update README.md

README.md

@@ -32,8 +32,7 @@ from transformers import AutoModelForMaskedLM, AutoModelForSequenceClassificatio
 import torch
 
 tokenizer  = AutoTokenizer.from_pretrained("biomap-research/xtrimopglm-1b-mlm", trust_remote_code=True, use_fast=True)
-config = AutoConfig.from_pretrained("biomap-research/xtrimopglm-1b-mlm",  trust_remote_code=True, torch_dtype=torch.bfloat16)
-model = AutoModelForMaskedLM.from_config(config, trust_remote_code=True, torch_dtype=torch.bfloat16)
+model = AutoModelForMaskedLM.from_pretrained("biomap-research/xtrimopglm-1b-mlm", trust_remote_code=True, torch_dtype=torch.bfloat16)
 if torch.cuda.is_available():
     model = model.cuda()
 model.eval()
@@ -46,10 +45,10 @@ with torch.inference_mode():
 
 
 # model for the sequence-level tasks
-model = AutoModelForSequenceClassification.from_config(config, trust_remote_code=True, torch_dtype=torch.bfloat16)
+model = AutoModelForSequenceClassification.from_pretrained("biomap-research/xtrimopglm-1b-mlm", trust_remote_code=True, torch_dtype=torch.bfloat16)
 
 # model for the token-level tasks
-model = AutoModelForTokenClassification.from_config(config, trust_remote_code=True, torch_dtype=torch.bfloat16)
+model = AutoModelForTokenClassification.from_pretrained("biomap-research/xtrimopglm-1b-mlm", trust_remote_code=True, torch_dtype=torch.bfloat16)
 
 ```