Views
No views yet
pip install git+https://github.com/huggingface/transformers.git1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3MODEL_PATH = "THUDM/glm-edge-4b-chat"
4
5tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
6model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto")
7
8message = [{"role": "user", "content": "hello!"}]
9
10inputs = tokenizer.apply_chat_template(
11 message,
12 return_tensors="pt",
13 add_generation_prompt=True,
14 return_dict=True,
15).to(model.device)
16
17generate_kwargs = {
18 "input_ids": inputs["input_ids"],
19 "attention_mask": inputs["attention_mask"],
20 "max_new_tokens": 128,
21 "do_sample": False,
22}
23out = model.generate(**generate_kwargs)
24print(tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))
25