Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_name = "DataPilot/Arrival-32B-Instruct-v0.3"
4tokenizer_name = ""
5
6if tokenizer_name == "":
7 tokenizer_name = model_name
8
9model = AutoModelForCausalLM.from_pretrained(
10 model_name,
11 torch_dtype="auto",
12 device_map="auto"
13)
14tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
15
16prompt = "9.9と9.11はどちらのほうが大きいですか？"
17messages = [
18 {"role": "system", "content": "あなたは優秀な日本語アシスタントです。問題解決をするために考えた上で回答を行ってください。"},
19 {"role": "user", "content": prompt}
20]
21text = tokenizer.apply_chat_template(
22 messages,
23 tokenize=False,
24 add_generation_prompt=True
25)
26model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
27
28generated_ids = model.generate(
29 **model_inputs,
30 max_new_tokens=1024
31)
32generated_ids = [
33 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
34]
35
36response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
37
38print(response)