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 f'{input}\n'
7 '\n'
8 '### 回答'
9 )
10
11 tokenized_input = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
12 attention_mask = torch.ones_like(tokenized_input)
13
14 with torch.no_grad():
15 outputs = model.generate(
16 tokenized_input,
17 attention_mask=attention_mask,
18 max_new_tokens=100,
19 do_sample=False,
20 repetition_penalty=1.2,
21 pad_token_id=tokenizer.eos_token_id
22 )[0]
23 output = tokenizer.decode(outputs[tokenized_input.size(1):], skip_special_tokens=True)
24 results.append({"task_id": data["task_id"], "input": input, "output": output})