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
11# Hugging Faceで取得したTokenをこちらに貼る。
12HF_TOKEN = userdata.get('HF_TOKEN')
13
14model_id = "llm-jp/llm-jp-3-13b"
15adapter_id = "tscp/llm-jp-3-13b-finetune"
16
17# QLoRA config
18bnb_config = BitsAndBytesConfig(
19 load_in_4bit=True,
20 bnb_4bit_quant_type="nf4",
21 bnb_4bit_compute_dtype=torch.bfloat16,
22)
23
24# Load model
25model = AutoModelForCausalLM.from_pretrained(
26 model_id,
27 quantization_config=bnb_config,
28 device_map="auto",
29 token = HF_TOKEN
30)
31
32# Load tokenizer
33tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, token = HF_TOKEN)
34
35# 元のモデルにLoRAのアダプタを統合。
36model = PeftModel.from_pretrained(model, adapter_id, token = HF_TOKEN)
37
38# データセットの読み込み。
39# (評価データセットのjsonlファイルのパスを設定してください)
40datasets = []
41with open("./elyza-tasks-100-TV_0.jsonl", "r") as f:
42 item = ""
43 for line in f:
44 line = line.strip()
45 item += line
46 if item.endswith("}"):
47 datasets.append(json.loads(item))
48 item = ""
49
50# gemma
51results = []
52for data in tqdm(datasets):
53 input = data["input"]
54 prompt = f"""### 指示
55 {input}
56 ### 回答
57 """
58
59 # input_ids だけを取り出して使用
60 input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
61 outputs = model.generate(input_ids, max_new_tokens=512, do_sample=False, repetition_penalty=1.2)
62 output = tokenizer.decode(outputs[0][input_ids.size(1):], skip_special_tokens=True)
63
64 results.append({"task_id": data["task_id"], "input": input, "output": output})
65
66# jsonl
67import re
68jsonl_id = re.sub(".*/", "", adapter_id)
69with open(f"./{jsonl_id}-outputs.jsonl", 'w', encoding='utf-8') as f:
70 for result in results:
71 json.dump(result, f, ensure_ascii=False) # ensure_ascii=False for handling non-ASCII characters
72 f.write('\n')