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 = "your-token"
13
14#ベースモデル (llm-jp-3-13b)
15base_model_id = "llm-jp/llm-jp-3-13b"
16#アダプタ
17adapter_id = "mamosaita/llm-jp-3-13b-finetune"
18
19# QLoRA config
20bnb_config = BitsAndBytesConfig(
21 load_in_4bit=True,
22 bnb_4bit_quant_type="nf4",
23 bnb_4bit_compute_dtype=torch.bfloat16,
24)
25
26# Load model
27model = AutoModelForCausalLM.from_pretrained(
28 model_id,
29 quantization_config=bnb_config,
30 device_map="auto",
31 token = HF_TOKEN
32)
33
34# Load tokenizer
35tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, token = HF_TOKEN)
36
37# 元のモデルにLoRAのアダプタを統合。
38model = PeftModel.from_pretrained(model, adapter_id, token = HF_TOKEN)
39
40# データセット(Elyza-tasks-100-TV_0.jsonl)の読み込み。
41datasets = []
42with open("./elyza-tasks-100-TV_0.jsonl", "r") as f:
43 item = ""
44 for line in f:
45 line = line.strip()
46 item += line
47 if item.endswith("}"):
48 datasets.append(json.loads(item))
49 item = ""
50
51# 推論
52results = []
53for data in tqdm(datasets):
54
55 input = data["input"]
56
57 prompt = f"""### 指示
58 {input}
59 ### 回答
60 """
61
62 tokenized_input = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
63 attention_mask = torch.ones_like(tokenized_input)
64 with torch.no_grad():
65 outputs = model.generate(
66 tokenized_input,
67 attention_mask=attention_mask,
68 max_new_tokens=100,
69 do_sample=False,
70 repetition_penalty=1.2,
71 pad_token_id=tokenizer.eos_token_id
72 )[0]
73 output = tokenizer.decode(outputs[tokenized_input.size(1):], skip_special_tokens=True)
74
75 results.append({"task_id": data["task_id"], "input": input, "output": output})