Update README.md

README.md

@@ -112,10 +112,17 @@ with torch.no_grad():
 from transformers import AutoTokenizer
 from vllm import LLM, SamplingParams
 
+# GLM-4-9B-Chat-1M
+# max_model_len, tp_size = 1048576, 4
+
 # GLM-4-9B-Chat
+from transformers import AutoTokenizer
+from vllm import LLM, SamplingParams
+
+# 如果遇见 OOM 现象，建议减少max_model_len，或者增加tp_size
 max_model_len, tp_size = 131072, 1
 model_name = "THUDM/glm-4-9b-chat"
-prompt = '你好'
+prompt = [{"role": "user", "content": "你好"}]
 
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 llm = LLM(
@@ -124,15 +131,17 @@ llm = LLM(
     max_model_len=max_model_len,
     trust_remote_code=True,
     enforce_eager=True,
+    # GLM-4-9B-Chat-1M 如果遇见 OOM 现象，建议开启下述参数
+    # enable_chunked_prefill=True,
+    # max_num_batched_tokens=8192
 )
 stop_token_ids = [151329, 151336, 151338]
 sampling_params = SamplingParams(temperature=0.95, max_tokens=1024, stop_token_ids=stop_token_ids)
 
-inputs = tokenizer.build_chat_input(prompt, history=None, role='user')['input_ids'].tolist()
-outputs = llm.generate(prompt_token_ids=inputs, sampling_params=sampling_params)
+inputs = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True)
+outputs = llm.generate(prompts=inputs, sampling_params=sampling_params)
 
-generated_text = [output.outputs[0].text for output in outputs]
-print(generated_text)
+print(outputs[0].outputs[0].text)
 ```
 
 ## 协议