Update README.md

README.md

@@ -25,8 +25,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 import torch
 
 tokenizer  = AutoTokenizer.from_pretrained("biomap-research/xtrimopglm-1b-clm", trust_remote_code=True, use_fast=True)
-config = AutoConfig.from_pretrained("biomap-research/xtrimopglm-1b-clm",  trust_remote_code=True, torch_dtype=torch.bfloat16)
-model = AutoModelForCausalLM.from_pretrained(config, trust_remote_code=True, torch_dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained("biomap-research/xtrimopglm-1b-clm", trust_remote_code=True, torch_dtype=torch.bfloat16)
 if torch.cuda.is_available():
     model = model.cuda()
 model.eval()
@@ -36,7 +35,7 @@ prompt=['', 'MLFVVL', 'LDL', 'VTQA']
 
 for idx, each in enumerate(prompt):
     print(f"Begin generating idx: {idx} with prompt {each}")
-    output = model.chat(tokenizer, each)
+    output = model.chat(tokenizer, each, **gen_kwargs)
     print(f"\nEnd generation with length: {len(output.split())} - seqs: {output}\n")
 ```