Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
2from peft import PeftModel, LoraConfig
3import torch
4from tqdm import tqdm
5
6HF_TOKEN = "Your HF token"
7BASE_MODEL_ID = "llm-jp/llm-jp-3-13b"
8ADAPTER_ID = "Enverop3/llm-jp-3-13b-finetune"
9
10# QLoRAの設定
11bnb_config = BitsAndBytesConfig(
12 load_in_4bit=True,
13 bnb_4bit_quant_type="nf4",
14 bnb_4bit_compute_dtype=torch.bfloat16,
15)
16
17# モデル読み込み
18model = AutoModelForCausalLM.from_pretrained(
19 BASE_MODEL_ID,
20 quantization_config=bnb_config,
21 device_map="auto",
22 token = HF_TOKEN
23)
24
25# Tokenizer読み込み
26tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True, token=HF_TOKEN)
27
28# FTしたアダプタを接続
29model = PeftModel.from_pretrained(model, ADAPTER_ID, token=HF_TOKEN)1results = []
2for data in tqdm(datasets):
3 input = data["input"]
4 prompt = (
5 '以下の指示文を確認し、以下のステップで適切な回答を出力します。\n'
6 '\n'
7 '[ステップ1]\n'
8 '適切な回答を導き出すために、指示文内で問われている内容は何か、ポイントを整理します。\n'
9 '整理したポイントは出力には含めず、以下のステップ2の回答のみを出力してください。\n'
10 '\n'
11 '[ステップ2]\n'
12 'ステップ1で整理したポイントに対して、適切な回答を出力します。\n'
13 '\n'
14 '### 指示文\n'
15 f'{input}\n'
16 '\n'
17 '### 回答'
18 )
19
20 tokenized_input = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
21 attention_mask = torch.ones_like(tokenized_input)
22
23 with torch.no_grad():
24 outputs = model.generate(
25 tokenized_input,
26 attention_mask=attention_mask,
27 max_new_tokens=100,
28 do_sample=False,
29 repetition_penalty=1.2,
30 pad_token_id=tokenizer.eos_token_id
31 )[0]
32 output = tokenizer.decode(outputs[tokenized_input.size(1):], skip_special_tokens=True)
33 results.append({"task_id": data["task_id"], "input": input, "output": output})