大規模言語モデル RWKV-worldで学習で巨大なデータ(学会の予稿集のpdf)をファインチューニング(LoRA)する

概要

  • 学会の予稿集のような、大規模な文章(pdf)データを大規模言語モデルに学習させてみます

    • 1.5 M tokenほどあります

  • モデルは、学習と出力が高速なRWKVにしました

    • 他のタスクでGPUリソースを使っているので、0.1B/ 1.5Bモデルでのお試しです

      • 1.5Bは学習時にVRAM 7GBほど使います

      • 執筆時、日本語最強のオープンLLMと謳われるRWKV-4-World-JPNtunedが本命ですが、7bを動かせるGPUが空いていなかったので、小さいモデルで試しています


前提: pdfデータの処理

LLMとは直接関係がありませんが、一般論として、テキスト学習にはデータの前処理が必要です。
今回は、数百MBのpdfデータとして存在する学会の予稿集をきれいなテキストに変換しました。
まずは、巨大すぎるpdfを分割しておきます。

#pdfの分割

import os
from PyPDF2 import PdfReader, PdfWriter
from tqdm import tqdm

def split_pdf(input_path, output_dir, chunk_size):
    with open(input_path, 'rb') as file:
        pdf = PdfReader(file)
        total_pages = len(pdf.pages)

        # ページを分割して新しいPDFファイルに保存する
        for i in tqdm(range(0, total_pages, chunk_size)):
            chunk_output = os.path.join(output_dir, f'output_{i+1}-{i+chunk_size}.pdf')
            with open(chunk_output, 'wb') as chunk_file:
                writer = PdfWriter()
                for j in range(i, min(i+chunk_size, total_pages)):
                    writer.add_page(pdf.pages[j])
                writer.write(chunk_file)
            print(f'Saved {chunk_output}')

# 分割するPDFファイルのパス (自分で決める)
input_pdf_path = "data/N72.pdf"

# 分割されたPDFファイルを保存するディレクトリ  (自分で決める)
output_directory = 'data/split_pdf'

# ページごとの分割サイズ
chunk_size = 10

# 出力ディレクトリが存在しない場合は作成する
if not os.path.exists(output_directory):
    os.makedirs(output_directory)

# PDFを分割する
split_pdf(input_pdf_path, output_directory, chunk_size)

pdfをテキストに変換します

from pdfminer.pdfinterp import PDFResourceManager, PDFPageInterpreter
from pdfminer.converter import TextConverter
from pdfminer.layout import LAParams
from io import StringIO
from pdfminer.pdfpage import PDFPage
import glob
import re

def pdf_to_text2(pdf_path):
    output_string = StringIO()
    with open(pdf_path, 'rb') as file:
        resource_manager = PDFResourceManager()
        converter = TextConverter(resource_manager, output_string, laparams=LAParams())
        interpreter = PDFPageInterpreter(resource_manager, converter)
        for page in (PDFPage.get_pages(file, check_extractable=True)):
            interpreter.process_page(page)
        converter.close()
        text = output_string.getvalue()
        output_string.close()
        return text


pdf_dir="data/split_pdf/*.pdf"
pdf_files=glob.glob(pdf_dir)


def extract_number(filename):
    filename = os.path.basename(filename)
    match = re.search(r'\d+', filename)
    return int(match.group()) if match else -1

pdf_files.sort(key=extract_number)

text=""
for idx in tqdm(range(len(pdf_files))):
    path=pdf_files[idx]
    t=pdf_to_text2(path)
    text+=t

テキストのクリーニング

pdfから抽出されたテキストには、学習に不向きな箇所が多々あります。例えば…

  • 変なところに挿入される改行コード

    • pdfのレイアウトの都合で挿入されます

  • FigureやTableの説明文

    • 今回はテキストだけを抜き出しているので、本文の途中に挿入される図表の説明文が、変なところに出現します

pdfから抽出したテキストのクリーニングはとても大変なのですが、今回は最小限のコードで対応します。

from tqdm import tqdm
import re
def has_alnum_or_symbol_start(string, n=5):
    start = string[:n]
    return bool(re.match(r'^[a-zA-Z0-9\W]+$', start))

page_split_tag="\x0c"


all_text=""

#ページで分割
abst_list=text.split(page_split_tag)
for idx in tqdm(range(len(abst_list))):
    abst=abst_list[idx]

    lines=abst.split("\n")

    lines=[line for line in lines if not has_alnum_or_symbol_start(line)]
    lines=[line for line in lines if not line.startswith(" ")]
    lines=[line for line in lines if line!=""]

    cleaned_text="".join(lines)
    all_text+=cleaned_text.strip()+"\n"

all_text=all_text.replace("\n\n","\n").strip()


#1行の文字数が30文字以下のものを削除
all_text="\n".join([line for line in all_text.split("\n") if len(line)>30])


#テキストを保存
with open(input_pdf_path.replace(".pdf",".txt"), "w") as f:
    f.write(all_text)

[補足説明]
has_alnum_or_symbol_start関数は、与えられた最初のn文字が、半角文字であるかを判定します。
例えば、「Fig. 1 …」で始まる行に対して、Trueを返し、「この論文では…」のような全角文字の文章に対しては、Falseを返します。

今回は日本語の予稿集を用いており、かつ、図表の表記はFig, Tableという感じだったので、この処理で不要な行をほとんど除くことができました(Referencesセクションも削除されます)。
英語の論文を扱う場合は、他のアルゴリズムを考える必要があります。

著作権の関係で、表示できないのですが、とりあえず、数MBのテキストデータが得られました。

本論: RWKVでの学習

RWKVはわりと盛り上がっており、各言語対応のRWKV-worldや、日本語最強と謳われるモデルが7月に出ました。
ファインチューニング関連は以下のページなどを参考にしました。

まずはトークナイズします。
上記の記事のコードでは、rwkv-world系のmodelやtokenizerに対応していないので、コードを書き換えていく必要があります

import numpy as np

from rwkv.utils import PIPELINE
from rwkv.model import RWKV

 #rwkv worldのtokenicer
model = RWKV(model='../RWKV-4-World-0.1B-v1-20230520-ctx4096', strategy='cpu fp32')
pipeline = PIPELINE(model, "rwkv_vocab_v20230424")

input_file="data/N72.txt"
output_file = 'train.npy'

with open(input_file,"r", errors='ignore') as f:
  txt=f.read()

#行末にスペシャルトークンを追加
s_token="<|endoftext|>"
lines=txt.split("\n")
lines=[i+s_token for i in lines]
txt="\n".join(lines)


data_code = pipeline.tokenizer.encode(txt)
print(f'Tokenized length = {len(data_code)}')

out = np.array(data_code, dtype='uint16')
np.save(output_file, out, allow_pickle=False)

ポイント

  • rwkv_vocab_v20230424というタイプのtokenizerが必要なので、RWKVのpipelineとして読み込んでいます。

    • この時点でmodelを読み込む必要はないのですが、他に良い手段が見つかりませんでした

  • 今回はQ&A形式ではなく、テキストをただそのまま読ませています

  • やっていることとしては、各文章の最後に、special tokenである<|endoftext|>を追加し、tokenizeしているだけです

Finetuning実行

RWKV-LM-LoRA/RWKV-v4neoのtrain.pyを使います(必要に応じてgitします)。

0.1bの場合


import os
os.environ['I_KNOW_WHAT_IM_DOING'] = 'True'

!cd  RWKV-LM-LoRA/RWKV-v4neo/ && python train.py \
  --load_model ../../../RWKV-4-World-0.1B-v1-20230520-ctx4096.pth \
  --proj_dir ../../out_model \
  --data_file "../../train.npy" \
  --data_type "numpy" \
  --vocab_size 65536  \
  --ctx_len 1024 \
  --epoch_save 5 \
  --epoch_count 100 \
  --n_layer 12 \
  --n_embd 768 \
  --epoch_steps 1000 --epoch_begin 0  --micro_bsz 1 --pre_ffn 0 --head_qk 0 --lr_init 1e-5\
  --lr_final 1e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.999 --adam_eps 1e-8\
  --accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_2 --grad_cp 0 \
  --lora --lora_r 32 --lora_alpha 32 --lora_dropout 0.1

1.5bの場合

#1.5b

import os
os.environ['I_KNOW_WHAT_IM_DOING'] = 'True'

!cd  RWKV-LM-LoRA/RWKV-v4neo/ && python train.py \
  --load_model ../../../RWKV-4-World-1.5B-v1-fixed-20230612-ctx4096.pth \
  --proj_dir ../../out_model \
  --data_file "../../train.npy" \
  --data_type "numpy" \
  --vocab_size 65536  \
  --ctx_len 1024 \
  --epoch_save 5 \
  --epoch_count 100 \
  --n_layer 24 \
  --n_embd 2048 \
  --epoch_steps 1000 --epoch_begin 0  --micro_bsz 1 --pre_ffn 0 --head_qk 0 --lr_init 1e-5\
  --lr_final 1e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.999 --adam_eps 1e-8\
  --accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_2 --grad_cp 0 \
  --lora --lora_r 32 --lora_alpha 32 --lora_dropout 0.1

ポイント・注意

  • load_modelにはベースとなるモデル、proj_dirにはloraモデルの出力先を指定します。

  • vocab_sizeは65536だと思いますが、少し自信ないです。

  • n_layerはmodel = RWKV(..)で読み込んだときに出てくる値から推測しています。

  • n_embedも、適当に打ち込んでエラーメッセージから推測してました。間違っているかもしれません

  • 'I_KNOW_WHAT_IM_DOING'を環境変数に設定しないと、バッチサイズが小さすぎると怒られます(errorになります)

  • 指定のepoch数を過ぎても、学習が無限に続くとの噂です。google cloab+などでは無駄にリソースを使うことになるので注意。

1.5bモデルの場合、RTX 2080で1 epochあたり、20分ほどかかります (1.5 M tokens)

LoRAモデルでの推論

RWKVのpipパッケージでは、LoRAモデルを読み込む機能が実装されていなかったので、gitをもとに実装します(このgitでオススメされているchat.pyを介したモデル読み込みでは、forward時にエラーが出てしまいました)。

パラメータ関連の設定(※関係のないパラメータがたくさん入ってます)


import types
import os
import sys
sys.path.append("RWKV-LM-LoRA/RWKV-v4neo/")
args = types.SimpleNamespace()
args.RUN_DEVICE = "cuda"  # 'cpu' (already very fast) // 'cuda'
args.FLOAT_MODE = "fp32" # fp32 (good for CPU) // fp16 (recommended for GPU) // bf16 (less accurate)
args.vocab_size = 50277
args.head_qk = 0
args.pre_ffn = 0
args.grad_cp = 0
args.my_pos_emb = 0

args.MODEL_NAME = '../RWKV-4-World-0.1B-v1-20230520-ctx4096'
args.MODEL_NAME="../RWKV-4-World-1.5B-v1-fixed-20230612-ctx4096"
args.n_layer = 13
args.n_embd =  768
args.ctx_len = 1024

# Modify this to u""se LoRA models; lora_r = 0 will not use LoRA weights.
args.MODEL_LORA = "out_model/rwkv-20"
args.lora_r = 32
args.lora_alpha = 32

os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '0'
os.environ["RWKV_T_MAX"]="1024"
os.environ["RWKV_FLOAT_MODE"]="fp32"
os.environ["RWKV_RUN_DEVICE"]="cuda"

ベースモデルの読み込み

fp16にすると、推論時にnanエラーが出やすい印象です

from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS

# モデルとパイプラインの準備
model = RWKV(
    model="../RWKV-4-World-1.5B-v1-fixed-20230612-ctx4096", 
    #model="../RWKV-4-World-0.1B-v1-20230520-ctx4096", 
    strategy="cuda fp32"
)

pipeline = PIPELINE(model, "rwkv_vocab_v20230424")

LoRAの重みの追加
※コードが間違っているかもしれません。ちょっと自信ないです。

#Loraの重みを追加
import torch
DEBUG_TIME=False
RWKV_RESCALE_LAYER = 6  # set x=x/2 every X layer
w=model.w
#merge
if args.lora_r > 0:
    # merge LoRA-only slim checkpoint into the main weights
    w_lora = torch.load(
        args.MODEL_LORA + '.pth', map_location='cpu')
    for k in w_lora.keys():
        w[k] = w_lora[k]
    # merge LoRA weights
    keys = set(w.keys())
    for k in keys:
        print(k)
        k: str
        if k.endswith('.weight'):
            prefix = k[:-len('.weight')]
            lora_A = prefix + '.lora_A'
            lora_B = prefix + '.lora_B'
            if lora_A in keys:
                assert lora_B in keys
                print(f'merging {lora_A} and {lora_B} into {k}')
                assert w[lora_B].shape[1] == w[lora_A].shape[0] == args.lora_r
                # merging needs matmul, which is slow on cpu; work on gpu if possible
                if args.RUN_DEVICE == 'cuda':
                    w[k] = w[k].cuda()
                    w[lora_A] = w[lora_A].cuda()
                    w[lora_B] = w[lora_B].cuda()
                w[k] += w[lora_B] @ w[lora_A] * \
                    (args.lora_alpha / args.lora_r)
                del w[lora_A]
                del w[lora_B]
# refine weights and send to correct device
keys = list(w.keys())
if 'pos_emb_x' in keys:
    w['pos_emb'] = (w['pos_emb_x'] + w['pos_emb_y']
                    ).reshape(args.ctx_len+1, -1)[:-1, :]
keys = list(w.keys())
print_need_newline = False
for x in keys:
    block_id = 0
    if 'blocks.' in x:
        block_id = int(x.split('.')[1])
    if 'att.output.weight' in x:
        w[x] = w[x] / (2 ** int(block_id // RWKV_RESCALE_LAYER))
    if 'ffn.value.weight' in x:
        w[x] = w[x] / (2 ** int(block_id // RWKV_RESCALE_LAYER))

    if '.time_' in x:
        w[x] = w[x].squeeze()
        if DEBUG_TIME:
            print(x, w[x].numpy())
    if '.time_decay' in x:
        w[x] = w[x].float()
        w[x] = -torch.exp(w[x])
    elif '.time_first' in x:
        w[x] = w[x].float()
    else:
        if args.FLOAT_MODE == "fp32":
            w[x] = w[x].float()
        elif args.FLOAT_MODE == "bf16":
            w[x] = w[x].bfloat16()
        elif args.FLOAT_MODE == "fp16":
            w[x] = w[x].half()

    w[x].requires_grad = False
    if args.RUN_DEVICE == 'cuda' and x != 'emb.weight':
        w[x] = w[x].cuda()

    if ('blocks.' not in x) or ('blocks.0.' in x):
        if print_need_newline:
            print('\n', end='')
            print_need_newline = False
        print(x.ljust(40), str(w[x].dtype).replace(
            'torch.', '').ljust(10), w[x].device)
    else:
        print_need_newline = True
        print('.', end='', flush=True)

推論

# Instructプロンプトの生成
def generate_prompt(instruction, input=None):
    if input:
        return f"""Instruction: {instruction}

Input: {input}

Response: """
    else:
        return f"""Question: {instruction}

Answer: """
# プロンプトの準備
prompt = generate_prompt(
    "高分子による熱伝導について教えて",
    "")
#print("--[prompt]--\n" + prompt + "----")

# 推論の実行
args = PIPELINE_ARGS(
    temperature = 1.0,
    top_p = 0.9, 
    top_k = 100, 
    alpha_frequency = 0.25, 
    alpha_presence = 0.25, 
    token_ban = [],
    token_stop = [0],
    chunk_len = 256) 
result = pipeline.generate(prompt, token_count=200, args=args)
print(result)

出力結果は以下の通り。

LoRA前 (0.1b)

高分子を含む静脈血圧による熱伝導変位のために生成された高分子物質は、細胞内の分泌物を調節する特性が特徴的である。 一般的な常連系統学で検討されている、新しい細胞内に設計されたプロヴァンスに類似するバイオマス酸解析菌原体(ABC)と呼ばれる不臓毒素検出方法が研究され、多くの存在を持つが、大量生成さえできない生理遺伝子となっている。これらの研究では、生成されるトリニトリケラトラナン等の糖を活性化させるための学習技術も用いられている。 Question: 小なりに似た現象は何かわかるように思われる意

LoRA10エポック (0.1b)
熱伝導論で、熱伝導線を生成する温度には、熱が進入する上で固体である。これは低温と表面分子の接続に必須となる表面熱条件である。熱伝導線は極上から熱に与える熱交換によって発生した場合、物質の層状态、環境空気、原子の密度によって密度変化が持ち、水溶液的な平衡により地表に浸入し暖かい熱伝導線が形成される。

7/18追記
LoRA35エポック(1.5b)
(loraの重み追加時に8GBではエラーが出たのでcpuで回しました。GPUでも頑張れば動くはずです)

高分子は、熱伝導率が非常に高く、あるいは体積密度の大きな単位円盤状の結晶を使用することで高い熱伝導性を発揮します。さらに、熱伝導率は非常に低く、一般的には 10-100 m2/W 未満となっていることが多くあります。


まとめ

  • RWKV-world系をLoRAしました

  • 0.1b~1.5bは元のモデルが間抜けなので、LoRA後もおかしな出力ですが、少しは良くなった気がします。

  • Q&A形式ではない状態でデータセットだけを学習させても、モデルが意外とQ&A形式で回答をしてくれた(?)のが、大きな収穫でした


この記事が気に入ったらサポートをしてみませんか?