日本語版LIMAデータセットでopen-calm-7bをLoRAファインチューニングしてチャットボットを作成してみた

!pip install peft
!pip install transformers
!pip install datasets
!pip install accelerate
!pip install bitsandbytes

accelerate==0.21.0.dev0
bitsandbytes==0.39.0
datasets==2.11.0
peft==0.4.0.dev0
transformers==4.27.4

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "cyberagent/open-calm-7b"
# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(model_name)
# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,
    device_map="auto",
    torch_dtype=torch.float16
)
# 複数のGPUを使う場合
if torch.cuda.device_count() > 1:
    model.is_parallelizable = True
    model.model_parallel = True

from datasets import load_dataset
dataset = load_dataset('zan/lima-ja')

>>> dataset["train"][0]
{'conversations': ['脳細胞は移動することができるのでしょうか？ 移動とは、遠距離の移動（可能ならば脳内のみで）を指します。',
  'この質問は比較的広範であり、脳は単にニューロンだけでなく、グリア細胞（支持細胞）や前ミトーシスの神経幹細胞も含まれることを考慮する必要があります。さらに、重要な同僚科学者たちが指摘しているように、発生段階も非常に重要であり、発達中の胎児の脳は成人の脳とは非常に異なります。\n\nしかし、さまざまな出版物を精査した結果、この質問の答えは実際には非常に単純であることがわかりました。結論として、脳細胞は移動します。\n\n成人の脳では、グリア細胞が脳内を移動します（Klämbt、2009年）。グリア細胞は多くの機能に関与していますが、移動するグリア細胞の注目すべき例は、標的の軸索を見つけるために比較的長距離を移動し、それに巻きついて絶縁性のミエリン鞘を形成するオリゴデンドロサイトです（TsaiとMiller、2002年）。\n\n神経幹細胞は損傷に対応して長距離を移動します（Imitola et al.、2004年）し、特定の幹細胞の位置（例えば、海馬や脳室下帯）から他の領域に移動します（Clarke、2003年）。\n\nミトーシス後でありながら未分化のニューロンは、魚（Scott et al.、2012年）や哺乳類、および非人猿でも成人の脳内を移動することが示されています（Sawada et al.、2011年）。\n\n驚くことではありませんが、グリア細胞、幹細胞、およびニューロンは胎児の発生中も移動します。特に、末梢機能を果たすために運命づけられた分裂後のニューロンは、神経堤から標的の位置まで比較的長い距離を移動しなければなりません（Neuroscience、第2版、Neuronal Migration）。'],
 'source': 'stackexchange'}

>>> dataset
DatasetDict({
    train: Dataset({
        features: ['conversations', 'source'],
        num_rows: 1030
    })
    test: Dataset({
        features: ['conversations', 'source'],
        num_rows: 300
    })
})

data_train = []
for i, item in enumerate(dataset["train"]):
    token = 0
    for j, c in enumerate(item["conversations"]):
        token += len(tokenizer.tokenize(c))
    if token <= 2048:
        data_train.append(item)

>>> len(data_train)
1006

from torch.utils.data import Dataset
from tqdm import tqdm

class ConversationDataset(Dataset):
    def __init__(self, json_list, tokenizer):
        # ConversationDatasetの初期化関数
        # json_list: JSON形式のデータリスト
        # tokenizer: トークナイザー
        self.tokenizer = tokenizer
        self.features = []
        all_texts_list = []

        # json_list内の各アイテムに対して処理を実行
        for item in tqdm(json_list):
            input_list = item["conversations"][:-1]
            output_text = item["conversations"][-1]
            input_text = "あなたはchatbotです。userと会話を行ってください。\n"
            
            # 入力リストの要素ごとに処理を実行
            for i, item in enumerate(input_list):
                # インデックスが偶数の場合はuserの発言を追加
                if i % 2 == 0:
                    input_text += f"\n### user:\n{item}\n"
                else:
                    # インデックスが奇数の場合はchatbotの発言を追加
                    input_text += f"\n### chatbot:\n{item}\n"
                    
            # chatbotの発言の欄を追加
            input_text += "\n### chatbot:\n"
            
            # 全てのテキストを結合し、トークナイザで処理
            all_text = input_text + output_text + self.tokenizer.eos_token
            all_texts_list.append(all_text)
        
        # トークナイズされたテキストのリストを特徴量として保存
        self.features = [
            tokenizer(
                text, padding=False, truncation=True, max_length=2048
            ) for text in tqdm(all_texts_list)
        ]
            
    def __len__(self):
        # データセットのサイズを返す
        return len(self.features)
    
    def __getitem__(self, idx):
        # 指定したインデックスの特徴量を返す
        return self.features[idx]

# インスタンスを作成
train_dataset = ConversationDataset(data_train, tokenizer)

>>> train_dataset[0]
{'input_ids': [7764, 1674, 899, 50143, 358, 247, 43302, 254, 5060, 44981, 247, 186, 186, 39843, ... 0],
 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ... 1]}

from transformers import DataCollatorForLanguageModeling

collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

import torch.nn as nn

# モデルのパラメーターごとに、勾配の計算を無効にする
for param in model.parameters():
    param.requires_grad = False

    # パラメーターの次元が1次元の場合、データ型をtorch.float32に変換する
    if param.ndim == 1:
        param.data = param.data.to(torch.float32)

# モデルに勾配チェックポイントを有効化する
model.gradient_checkpointing_enable()

# モデルの入力に対する勾配の計算を有効化する
model.enable_input_require_grads()

# CastOutputToFloatという名前のクラスを定義し、nn.Sequentialクラスを継承する
class CastOutputToFloat(nn.Sequential):
    # forwardメソッドをオーバーライドし、入力データxに対してsuper()を呼び出して順伝搬を実行し、
    # 結果をtorch.float32のデータ型に変換して返す
    def forward(self, x):
        return super().forward(x).to(torch.float32)

# モデルのembed_out属性にCastOutputToFloatクラスのインスタンスを代入する
model.embed_out = CastOutputToFloat(model.embed_out)

from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, TaskType

# LoRAのパラメータ
lora_config = LoraConfig(
    r= 8, 
    lora_alpha=16,
    target_modules=["query_key_value", "dense", "dense_h_to_4h", "dense_4h_to_h"],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

# モデルの前処理
model = prepare_model_for_int8_training(model)

# LoRAモデルの準備
model = get_peft_model(model, lora_config)

# 学習可能パラメータの確認
model.print_trainable_parameters()

import transformers

peft_name = "lora-calm-7b"
output_dir = "lora-calm-7b-output"

# トレーナーの準備
trainer = transformers.Trainer(
    model=model,
    train_dataset=train_dataset,
    args=transformers.TrainingArguments(
        num_train_epochs=1,
        learning_rate=3e-4,
        logging_steps=20,
        output_dir=output_dir,
        report_to="none",
        save_total_limit=1,
        push_to_hub=False,
        auto_find_batch_size=True,
    ),
    data_collator=collator,
)

# 学習の実行
model.config.use_cache = False
trainer.train()
model.config.use_cache = True

# LoRAモデルの保存
trainer.model.save_pretrained(peft_name)

from peft import PeftModel

model_name = "cyberagent/open-calm-7b"
# トークンナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(model_name)
# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,
    device_map="auto",
)
peft_name = "./lora-calm-7b"
# LoRAモデルの準備
model = PeftModel.from_pretrained(
    model, 
    peft_name, 
    device_map="auto",
　　torch_dtype=torch.float16
)

# 評価モードへ
model.eval()

def generate_prompt(input_text):
    conversation_text = "あなたはchatbotです。userと会話を行ってください。\n"
    conversation_text += f"\n### user:\n{input_text}\n"
    conversation_text += "\n### chatbot:\n"
    return conversation_text

input_text=""

input_ids = tokenizer(generate_prompt(input_text), return_tensors="pt", truncation=True).input_ids.cuda()
outputs = model.generate(
    input_ids=input_ids,
    max_new_tokens=1024, 
    do_sample=True,
    temperature=0.6, 
    top_p=0.8, 
    top_k=40,
    no_repeat_ngram_size=2,
)
outputs = outputs[0].tolist()

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

    # レスポンス内容のみ抽出
    sentinel = "### chatbot:"
    sentinelLoc = decoded.find(sentinel)
    if sentinelLoc >= 0:
        print(decoded[sentinelLoc+len(sentinel):])
    else:
        print('Warning: Expected prompt template to be emitted.  Ignoring output.')
        print("------------------")
        print(tokenizer.decode(outputs))
else:
    print('Warning: no <eos> detected ignoring output')
    print("------------------")
    print(tokenizer.decode(outputs))