Views
No views yet

1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4device = "cuda"
5
6tokenizer = AutoTokenizer.from_pretrained("HODACHI/glm-4-9b-chat-FT-ja-v0.3", trust_remote_code=True)
7
8prompt = f"""
9リンゴとバナナはどのように似ていますか?
10"""
11query = "あなたは優秀なアシスタントです。質問を熟考して正確に答えるようにしてください。\n\n" + prompt
12
13inputs = tokenizer.apply_chat_template([{"role": "user", "content": query}],
14 add_generation_prompt=True,
15 tokenize=True,
16 return_tensors="pt",
17 return_dict=True
18 )
19
20inputs = inputs.to(device)
21model = AutoModelForCausalLM.from_pretrained(
22 "HODACHI/glm-4-9b-chat-FT-ja-v0.3",
23 torch_dtype=torch.bfloat16,
24 low_cpu_mem_usage=True,
25 trust_remote_code=True
26).to(device).eval()
27
28gen_kwargs = {"max_length": 2500, "do_sample": True, "top_k": 1}
29with torch.no_grad():
30 outputs = model.generate(**inputs, **gen_kwargs)
31 outputs = outputs[:, inputs['input_ids'].shape[1]:]
32 print(tokenizer.decode(outputs[0], skip_special_tokens=True))