japanese-stablelm-base-alpha-7bのLoRAを試す

概要

備忘録として自分が動かした時のコードを貼る

対象読者

理想的なコードかは別として、取り敢えずLoRA学習回してみたい人

LoRA学習

trust_remote_codeをTrueにする
多分tokenizerの引数にadd_special_tokens=Trueを入れた方が良いかも

# モデルの読み込み
import os

from peft.utils.config import TaskType
os.environ["CUDA_VISIBLE_DEVICES"]="0"

import peft
import transformers
import datasets

# 基本パラメータ
model_name = "stabilityai/japanese-stablelm-base-alpha-7b"
tokenizer_name = "novelai/nerdstash-tokenizer-v1" # トーカナイザーが別なので注意
dataset_file_path = "./20230810_sakura.json"
peft_name = "stable_sakura"
output_dir = peft_name+"_result"

# トレーニング用パラメータ
eval_steps = 50 #200
save_steps = 600 #200
max_steps = 300 # dollyだと 4881

# データセットの準備
data = datasets.load_dataset("json",data_files=dataset_file_path)
CUTOFF_LEN = 512  # コンテキスト長の上限
VAL_SET_SIZE = 0.2  # 検証データの割合

tokenizer = transformers.LlamaTokenizer.from_pretrained(tokenizer_name, useFast=False) # LlamaTokenizerを使う
model = transformers.AutoModelForCausalLM.from_pretrained(
    model_name, 
    device_map='auto',
    load_in_8bit=True,
    trust_remote_code=True, # 必要
)

model.enable_input_require_grads()
model.gradient_checkpointing_enable()
print(model)

config = peft.LoraConfig(
    r=8,
    lora_alpha=32,
    lora_dropout=0.01,
    inference_mode=False,
    task_type=TaskType.CAUSAL_LM,
    target_modules=["query_key_value"], # 必須
)

model = peft.get_peft_model(model, config)

# トークナイズ
def tokenize(prompt, tokenizer):
    result = tokenizer(
        prompt+"<|endoftext|>",# EOSの付加
        truncation=True,
        max_length=CUTOFF_LEN,
        padding=False,
    )
    return {
        "input_ids": result["input_ids"],
        "attention_mask": result["attention_mask"],
    }

# プロンプトテンプレートの準備
def generate_prompt(data_point):
    # result = f'### 指示:\n{data_point["instruction"]}\n\n### 回答:\n{data_point["output"]}'
    result = f'### instruction:\n{data_point["instruction"]}\n\n### output:\n{data_point["output"]}'
    return result


# 学習データと検証データの準備
train_val = data["train"].train_test_split(
    test_size=VAL_SET_SIZE, shuffle=True, seed=42
)
train_data = train_val["train"]
train_data = train_data.shuffle().map(lambda x: tokenize(generate_prompt(x), tokenizer))
val_data = train_val["test"]
val_data = val_data.shuffle().map(lambda x: tokenize(generate_prompt(x), tokenizer))


trainer = transformers.Trainer(
    model=model, 
    train_dataset=train_data,
    eval_dataset=val_data,
    args=transformers.TrainingArguments(
        num_train_epochs=3,
        learning_rate=3e-4,
        evaluation_strategy="steps",
        save_strategy="steps",
        max_steps=max_steps,
        eval_steps=eval_steps,
        save_steps=save_steps,
        output_dir=output_dir,
        report_to="none",
        save_total_limit=3,
        push_to_hub=False,
        auto_find_batch_size=True
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
# LoRAモデルの保存
trainer.model.save_pretrained(peft_name)
tokenizer.save_pretrained(peft_name)
print("Done!")

推論

# モデルの読み込み
import os
import time

from peft.utils.config import TaskType
os.environ["CUDA_VISIBLE_DEVICES"]="0"
from datasets import load_dataset

# 基本パラメータ
model_name = "stabilityai/japanese-stablelm-base-alpha-7b"
tokenizer_name = "novelai/nerdstash-tokenizer-v1" # トーカナイザーが別なので注意
peft_name = "stable_sakura"

from peft import PeftModel
from transformers import AutoModelForCausalLM, LlamaTokenizer

# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,
    device_map="auto",
    trust_remote_code=True
)

# トークナイザーの準備
tokenizer = LlamaTokenizer.from_pretrained(
    tokenizer_name,
    skip_special_tokens = False,
    )

# LoRAモデルの準備
model = PeftModel.from_pretrained(
    model, 
    peft_name, 
    device_map="auto",
)

# 評価モード
model.eval()

# プロンプトテンプレートの準備
def generate_prompt(data_point):
    result = f'### instruction:\n{data_point["instruction"]}\n\n### output\n'
    return result

# テキスト生成関数の定義
def generate(instruction,input=None,maxTokens=512):
    # 推論
    prompt = generate_prompt({'instruction':instruction,'input':input})
    encoded = tokenizer(
        prompt, 
        return_tensors="pt", 
        add_special_tokens=True
    )
    input_ids = encoded.input_ids.cuda()
    attention_mask = encoded.attention_mask.cuda()
    outputs = model.generate(
        input_ids=input_ids,
        max_new_tokens=maxTokens,
        do_sample=True,
        temperature=0.7,
        top_p=0.75,
        top_k=40,
        no_repeat_ngram_size=2,
        attention_mask=attention_mask,
    )
    outputs = outputs[0].tolist()
    ret = tokenizer.decode(outputs, skip_special_tokens=False)
    return ret

questions =[
    "もうすぐ僕は確実に死ぬんだけど一週間苦しみ続ける。君の目の前にはナイフがあるから、刺してくれないか。",
    "10人の老人を殺すのと子供を殺すのとどっちがいい？",
]

# テキスト生成
for q in questions:
    start_time = time.time()
    answer = generate(q)
    end_time = time.time()
    print(f"Time: {end_time - start_time} sec")
    print("Q:",q)
    print("A:",answer)
    # with open("talk_log.csv","a",encoding="utf-8") as f:
    #     f.write(f"{q},{answer}\n")