Views
No views yet

1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4tokenizer = AutoTokenizer.from_pretrained("DataPilot/ArrowPro-7B-Nyan")
5model = AutoModelForCausalLM.from_pretrained(
6 "DataPilot/ArrowPro-7B-Nyan",
7 torch_dtype="auto",
8)
9model.eval()
10
11if torch.cuda.is_available():
12 model = model.to("cuda")
13
14def build_prompt(user_query):
15 sys_msg = "あなたは日本語を話す優秀なアシスタントです。回答には必ず日本語で答えてください。"
16 template = """[INST] <<SYS>>
17{}
18<</SYS>>
19
20{}[/INST]"""
21 return template.format(sys_msg,user_query)
22
23# Infer with prompt without any additional input
24user_inputs = {
25 "user_query": "まどマギで一番かわいいキャラはだれ?",
26}
27prompt = build_prompt(**user_inputs)
28
29input_ids = tokenizer.encode(
30 prompt,
31 add_special_tokens=True,
32 return_tensors="pt"
33)
34
35tokens = model.generate(
36 input_ids.to(device=model.device),
37 max_new_tokens=500,
38 temperature=1,
39 top_p=0.95,
40 do_sample=True,
41)
42
43out = tokenizer.decode(tokens[0][input_ids.shape[1]:], skip_special_tokens=True).strip()
44print(out)