Views
No views yet
1from transformers import (
2 AutoModelForCausalLM,
3 AutoTokenizer,
4 BitsAndBytesConfig,
5)
6from peft import PeftModel
7import torch
8from tqdm import tqdm
9import json
10
11HF_TOKEN = 'YOUR TOKEN'
12model_id = "llm-jp/llm-jp-3-13b"
13adapter_id = "chanosuke/llm-jp-3-13b-finetune-3"
14
15# QLoRA config
16bnb_config = BitsAndBytesConfig(
17 load_in_4bit=True,
18 bnb_4bit_quant_type="nf4",
19 bnb_4bit_compute_dtype=torch.bfloat16,
20)
21
22# Load model
23model = AutoModelForCausalLM.from_pretrained(
24 model_id,
25 quantization_config=bnb_config,
26 device_map="auto",
27 token = HF_TOKEN
28)
29
30# Load tokenizer
31tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, token = HF_TOKEN)
32
33# 元のモデルにLoRAのアダプタを統合。
34model = PeftModel.from_pretrained(model, adapter_id, token = HF_TOKEN)
35
36# データセットの読み込み。
37# 推論環境内のelyza-tasks-100-TV_0.jsonalのパスを入力
38datasets = []
39with open('./elyza-tasks-100-TV_0.jsonl', "r") as f:
40 item = ""
41 for line in f:
42 line = line.strip()
43 item += line
44 if item.endswith("}"):
45 datasets.append(json.loads(item))
46 item = ""
47
48# llmjp推論
49# 学習したモデルを用いてタスクを実行
50from tqdm import tqdm
51
52# 推論するためにモデルのモードを変更
53model.eval()
54# モデルがGPUで動作する場合、GPUに移動
55device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
56model.to(device)
57
58results = []
59for dt in tqdm(datasets):
60 input=dt['input']
61 prompt = f"""
62### 指示
63以下の形式に従って回答を生成してください。
64
65入力:
66{input}
67
68### 回答
69"""
70
71 tokenized_input = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
72 with torch.no_grad():
73 outputs = model.generate(
74 tokenized_input,
75 max_new_tokens=1024,
76 #num_return_sequences=3,
77 do_sample=True,
78 top_p=0.95,
79 temperature=0.7,
80 repetition_penalty=1.05,
81 pad_token_id=tokenizer.eos_token_id
82 )
83 output_texts = [tokenizer.decode(output[tokenized_input.size(1):], skip_special_tokens=True) for output in outputs]
84 prediction = output_texts[0]
85 results.append({"task_id": dt["task_id"], "input": input, "output": prediction})
86
87
88# 推論結果をjsonlファイルに出力する
89import re
90jsonl_id = re.sub(".*/", "", adapter_id)
91with open(f"./{jsonl_id}-outputs.jsonl", 'w', encoding='utf-8') as f:
92 for result in results:
93 json.dump(result, f, ensure_ascii=False) # ensure_ascii=False for handling non-ASCII characters
94 f.write('\n')