Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。特に指示が無い場合は、常に日本語で回答してください。"
5text = "優秀なAIとはなんですか？ またあなたの考える優秀なAIに重要なポイントを５つ挙げて下さい。"
6
7model_name = "DataPilot/Llama3-ArrowSE-8B-v0.1"
8
9tokenizer = AutoTokenizer.from_pretrained(model_name)
10model = AutoModelForCausalLM.from_pretrained(
11 model_name,
12 torch_dtype="auto",
13 device_map="auto",
14)
15model.eval()
16
17messages = [
18 {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
19 {"role": "user", "content": text},
20]
21prompt = tokenizer.apply_chat_template(
22 messages,
23 tokenize=False,
24 add_generation_prompt=True
25)
26token_ids = tokenizer.encode(
27 prompt, add_special_tokens=False, return_tensors="pt"
28)
29
30with torch.no_grad():
31 output_ids = model.generate(
32 token_ids.to(model.device),
33 max_new_tokens=1200,
34 do_sample=True,
35 temperature=0.6,
36 top_p=0.9,
37 )
38output = tokenizer.decode(
39 output_ids.tolist()[0][token_ids.size(1):], skip_special_tokens=True
40)
41print(output)