Modify to original glm-4-9b code

README.md

@@ -28,20 +28,10 @@ tokenizer = AutoTokenizer.from_pretrained("THUDM/LongWriter-glm4-9b", trust_remo
 model = AutoModelForCausalLM.from_pretrained("THUDM/LongWriter-glm4-9b", torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto")
 model = model.eval()
 query = "Write a 10000-word China travel guide"
-prompt = f"[INST]{query}[/INST]"
-input = tokenizer(prompt, truncation=False, return_tensors="pt").to(device)
-context_length = input.input_ids.shape[-1]
-output = model.generate(
-    **input,
-    max_new_tokens=32768,
-    num_beams=1,
-    do_sample=True,
-    temperature=0.5,
-)[0]
-response = tokenizer.decode(output[context_length:], skip_special_tokens=True)
+response, history = model.chat(tokenizer, query, history=[], max_new_tokens=32768, temperature=0.5)
 print(response)
 ```
-Environment: `transformers==4.43.0`
+Environment: Same environment requirement as [glm-4-9b-chat](https://huggingface.co/THUDM/glm-4-9b-chat) (`transforemrs>=4.44.0`).
 
 License: [glm-4-9b License](https://huggingface.co/THUDM/glm-4-9b-chat/blob/main/LICENSE)