PEFTをWindowsから遊ぶメモ

conda create -n PEFT python=3.10
activate PEFT

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install datasets accelerate loralib sentencepiece transformers
python -m pip install https://github.com/jllllll/bitsandbytes-windows-webui/raw/main/bitsandbytes-0.39.0-py3-none-any.whl
pip install git+https://github.com/huggingface/peft.git

from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_int8_training
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from datasets import load_dataset

# 基本パラメータ
base_model = "rinna/japanese-gpt-neox-3.6b-instruction-ppo"
dataset = "bbz662bbz/databricks-dolly-15k-ja-gozaru"
peft_name = "test-rinna-3.6b"
output_dir = "test-rinna-3.6b-results"

# トレーニング用パラメータ
eval_steps = 100
save_steps = 200
logging_steps = 20
epochs = 3
max_steps = 250 # 0にするとepochsに応じて自動設定

# LoRA用パラメータ
lora_r = 8
lora_alpha = 16
lora_dropout = 0.1

# 他パラメータ
VAL_SET_SIZE = 0.2 # 検証分割比率
CUTOFF_LEN = 512  # コンテキスト長の上限

# ベースモデルの読み込み
model = AutoModelForCausalLM.from_pretrained(
    base_model, 
    device_map='auto',
    load_in_8bit=True,
)
tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=False) # Rinnaのトークナイザーでは、「use_fast=False」も必要になる

# PEFT(LoRA)の設定
config = LoraConfig(r=lora_r,
                    lora_alpha=lora_alpha,
                    lora_dropout=lora_dropout,
                    inference_mode=False,
                    task_type=TaskType.CAUSAL_LM,
)
model = prepare_model_for_int8_training(model)
model = get_peft_model(model, config)
model.enable_input_require_grads()  
model.gradient_checkpointing_enable()   
model.print_trainable_parameters()  # 学習可能パラメータの確認

# トークナイズ関数
def tokenize(prompt, tokenizer):
    result = tokenizer(prompt,
                       truncation=True,
                       max_length=CUTOFF_LEN,
                       padding=False,
    )
    return {"input_ids": result["input_ids"],
            "attention_mask": result["attention_mask"],
    }

# データセットの準備
data = load_dataset(dataset)

# プロンプトテンプレートの準備
# 回答の最後にベースモデル特有のEOSトークン</s>を挿入
def generate_prompt(data_point):
    if data_point["input"]:
        result = f"""ユーザー:{data_point["instruction"]}
入力:{data_point["input"]}
システム:{data_point["output"]}</s>"""
    else:
        result = f"""ユーザー:{data_point["instruction"]}
システム:{data_point["output"]}</s>"""

    result = result.replace('\n', '<NL>')   # 改行→<NL>
    return result

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

# トレーナーの定義
trainer = Trainer(
    model = model, 
    train_dataset = train_data,
    eval_dataset = val_data,
    args = TrainingArguments(
        num_train_epochs=epochs,
        learning_rate=3e-4,
        logging_steps=logging_steps,
        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= DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

model.config.use_cache = False  # 警告を黙らせます
trainer.train()
model.config.use_cache = True   # 推論のために再度有効化
trainer.model.save_pretrained(peft_name)    # LoRAモデルの保存
print("Done!")

import time
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# 基本パラメータ
base_model = "rinna/japanese-gpt-neox-3.6b-instruction-ppo"
peft_name = "test-rinna-3.6b"

# 入力文章
input_text = "カレーにジャガイモは入れるべき？"

# ベースモデルの読み込み
model = AutoModelForCausalLM.from_pretrained(
    base_model, 
    device_map='auto',
    load_in_8bit=True,
)
tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=False) # Rinnaのトークナイザーでは、「use_fast=False」も必要になる

# PEFT(LoRA)の読み込み
model = PeftModel.from_pretrained(model, peft_name)
model.eval()# 評価モード

# プロンプトテンプレートの準備
def generate_prompt(data_point):
    result = f"""ユーザー:{data_point["instruction"]}
システム:"""

    result = result.replace('\n', '<NL>')   # 改行→<NL>
    return result

# テキスト生成関数の定義
def textgen(instruction,input=None,maxTokens=256):
    prompt = generate_prompt({'instruction':instruction,'input':input})
    input_ids = tokenizer(prompt, 
                          return_tensors="pt", 
                          truncation=True, 
                          add_special_tokens=False).input_ids.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,
    )
    outputs = outputs[0].tolist()
    #print(tokenizer.decode(outputs))

    # EOSトークンにヒットしたらデコード完了
    if tokenizer.eos_token_id in outputs:
        eos_index = outputs.index(tokenizer.eos_token_id)
        decoded = tokenizer.decode(outputs[:eos_index])

        # レスポンス内容のみ抽出
        sentinel = "システム:"
        sentinelLoc = decoded.find(sentinel)
        if sentinelLoc >= 0:
            result = decoded[sentinelLoc+len(sentinel):]
            return(result.replace("<NL>", "\n"))  # <NL>→改行
        else:
            return('Warning: Expected prompt template to be emitted.  Ignoring output.')
    else:
        return('Warning: no <eos> detected ignoring output')

start_time = time.time()
print(f"Q: {input_text}\nA: {textgen(input_text)}")
end_time = time.time()
elapsed_time = end_time - start_time
print(f"Elapsed time: {elapsed_time} seconds")