テキストファイルからのLLM事前学習

今回は、スクラッチでLLMの事前学習を行う基本的な流れを確認するためにテキストファイルを読み込んでLLMで学習させてみることにします。

基本的には前回と同様ですが、数十の文を学習させるとどんな感じになるか、見てみましょう。

テキストファイルを扱うTokenizerとDataset

今回使用するテキストファイルは、以下のように単純な短文を36個並べたものにしました。このテキストファイルの各行を一つのデータとして扱います

道を歩く。
道を歩いた。
家まで歩く。
家まで歩いた。
駅から歩く。
駅から歩いた。
本を読む。
本を読んだ。
席に座る。
席に座った。
窓を開ける。
窓を開けた。
服を着る。
服を着た。
靴を履く。
靴を履いた。
私は、道を歩く。
私は、道を歩いた。
私は、家まで歩く。
私は、家まで歩いた。
私は、駅から歩く。
私は、駅から歩いた。
私は、本を読む。
私は、本を読んだ。
私は、席に座る。
私は、席に座った。
私は、窓を開ける。
私は、窓を開けた。
私は、服を着る。
私は、服を着た。
私は、靴を履く。
私は、靴を履いた。
私は、暗い道を歩く。
私は、暗い道を歩いた。
私は、駅から家まで歩く。
私は、駅から家まで歩いた。

この一連のデータ文を扱うトークナイザーは、以下の設定ファイルから作成します。

{
    "version": "1.0",
    "truncation": null,
    "padding": null,
    "normalizer": null,
    "pre_tokenizer": null,
    "post_processor": null,
    "model": {
        "type": "Unigram",
        "vocab": [
            ["<:pad:>", 0],
            ["<:eos:>", 1],
            ["、", 2],
            ["。", 3],
            ["道", 4],
            ["を", 5],
            ["歩", 6],
            ["く", 7],
            ["い", 8],
            ["た", 9],
            ["家", 10],
            ["ま", 11],
            ["で", 12],
            ["駅", 13],
            ["か", 14],
            ["ら", 15],
            ["本", 16],
            ["読", 17],
            ["む", 18],
            ["ん", 19],
            ["だ", 20],
            ["席", 21],
            ["に", 22],
            ["座", 23],
            ["る", 24],
            ["っ", 25],
            ["窓", 26],
            ["開", 27],
            ["け", 28],
            ["服", 29],
            ["着", 30],
            ["靴", 31],
            ["履", 32],
            ["私", 33],
            ["は", 34],
            ["暗", 35],
            ["い", 36]
        ]
    }
}

トークナイザーに合わせてLLMの設定ファイルにあるvocab_sizeも修正しておきます。

{
    "architectures": [
        "GPTNeoXForCausalLM"
    ],
    "hidden_act": "gelu",
    "hidden_size": 768,
    "initializer_range": 0.02,
    "intermediate_size": 3072,
    "layer_norm_eps": 1e-05,
    "max_position_embeddings": 2048,
    "model_type": "gpt_neox",
    "num_attention_heads": 12,
    "num_hidden_layers": 12,
    "rotary_emb_base": 10000,
    "rotary_pct": 1.0,
    "tie_word_embeddings": false,
    "torch_dtype": "float32",
    "transformers_version": "4.42.0",
    "use_cache": true,
    "use_parallel_residual": false,
    "vocab_size": 37
}

学習とテスト

データとトークナイザーが用意ができたので、ファイルを読み込んで1行ずつトークン化したデータを作成するDatasetと共に学習コードを作ってみましょう。

import torch
from torch.utils.data import Dataset
from transformers import AutoTokenizer, DataCollatorForLanguageModeling
from transformers import GPTNeoXConfig, GPTNeoXForCausalLM
from transformers import Trainer, TrainingArguments

class TestDataset(Dataset):

    def __init__(self, tokenizer):

        self.data_list = []

        with open("train.txt", 'r', encoding="utf-8") as f:

            lines = f.read().split('\n')

            for line in lines:

                input_ids = tokenizer.encode(line + "<:eos:>", return_tensors='pt')[0]
                self.data_list.append({'input_ids': input_ids})

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, idx):
        return self.data_list[idx]

model_config = GPTNeoXConfig.from_json_file("model_config_small.json")

model = GPTNeoXForCausalLM(model_config)
tokenizer = AutoTokenizer.from_pretrained("tokenizer")

dataset = TestDataset(tokenizer)

collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

training_args = TrainingArguments(output_dir="output", save_steps=36, num_train_epochs=1, learning_rate=5e-5, lr_scheduler_type="constant", per_device_train_batch_size=1)

trainer = Trainer(model=model, data_collator=collator, args=training_args, train_dataset=dataset)

trainer.train()

Datasetのデータ作成部分以外は特に変更はありません。

このコードで全データを1回だけ学習させてから、「私」「私は」「道」「窓を」「服」「靴を」「駅から」「駅から家」に続く語を出力してみます。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("output/checkpoint-36")
tokenizer = AutoTokenizer.from_pretrained("tokenizer")

input_texts = ["私", "私は", "道", "窓を", "服", "靴を", "駅から", "駅から家"]

for input_text in input_texts:

    input = tokenizer(input_text, return_tensors="pt", return_token_type_ids=False)
    texts = model.generate(**input, max_new_tokens=16, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, do_sample=True, num_return_sequences=4)

    for text in texts:
        print(tokenizer.decode(text))
結果は、以下のようになりました。

私 は 、 か で 歩 く 。 <:eos:> <:pad:> <:pad:>
私 は 、 る 。 <:eos:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:>
私 は 、 ま で 歩 く 。 <:eos:> <:pad:> <:pad:>
私 は 、 駅 ら 歩 歩 歩 く 。 <:eos:>
私 は 、 に 。 <:eos:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:>
私 は 、 に ら 。 <:eos:> <:pad:> <:pad:> <:pad:> <:pad:>
私 は 、 家 ま で 歩 く 。 <:eos:> <:pad:>
私 は 、 け 家 ま で 歩 く 。 <:eos:>
道 を い 。 <:eos:> <:pad:> <:pad:>
道 を く く い 。 <:eos:>
道 座 を 歩 。 <:eos:> <:pad:>
道 を 歩 く た 。 <:eos:>
窓 を 開 け け 開 る け 。 <:eos:> <:pad:>
窓 を る 歩 開 け 。 <:eos:> <:pad:> <:pad:> <:pad:>
窓 を け 開 る け け る た 。 <:eos:>
窓 を け る 開 け 開 け る 。 <:eos:>
服 を る る 着 る る 着 を け る る 着 。 <:eos:> <:pad:> <:pad:>
服 を 着 着 着 る 着 る 着 る る る る を る 。 <:eos:>
服 を 着 る る 着 着 る る 。 <:eos:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:>
服 を 着 る 着 る 着 る る る る か る 。 <:eos:> <:pad:> <:pad:>
靴 を 履 た 。 <:eos:>
靴 を い た 。 <:eos:>
靴 を い い 。 <:eos:>
靴 を た た 。 <:eos:>
駅 か ら 歩 い た 歩 い 。 <:eos:>
駅 か ら 歩 い た 。 <:eos:> <:pad:> <:pad:>
駅 か ら 歩 歩 た 。 <:eos:> <:pad:> <:pad:>
駅 か ら 歩 歩 く 。 <:eos:> <:pad:> <:pad:>
駅 か ら 家 。 <:eos:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:>
駅 か ら 家 歩 く く た 。 <:eos:> <:pad:> <:pad:>
駅 か ら 家 歩 ん 歩 い た た 。 <:eos:>
駅 か ら 家 歩 い く 。 <:eos:> <:pad:> <:pad:> <:pad:>

一部上手く続いている部分もありますが、データが多いこともあってかなり乱れていますね。4回ずつ学習させると…

私 は 、 窓 を 開 け た 。 <:eos:> <:pad:> <:pad:> <:pad:>
私 は 、 暗 い た 。 <:eos:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:>
私 は 、 席 に 座 道 を 歩 い た 。 <:eos:>
私 は 、 本 を 読 む 。 <:eos:> <:pad:> <:pad:> <:pad:> <:pad:>
私 は 、 家 ま で 歩 い た 。 <:eos:> <:pad:> <:pad:> <:pad:>
私 は 、 席 に 座 た 。 <:eos:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:>
私 は 、 道 を 歩 い た 。 <:eos:> <:pad:> <:pad:> <:pad:> <:pad:>
私 は 、 駅 か ら 家 ま で 歩 い た 。 <:eos:>
道 を 歩 い た 。 <:eos:>
道 を 歩 い た 。 <:eos:>
道 を 歩 い た 。 <:eos:>
道 を 歩 い た 。 <:eos:>
窓 を 開 け た 。 <:eos:>
窓 を 開 け た 。 <:eos:>
窓 を 開 け た 。 <:eos:>
窓 を 開 け た 。 <:eos:>
服 を 着 た 。 <:eos:>
服 を 着 た 。 <:eos:>
服 を 着 た 。 <:eos:>
服 を 着 た 。 <:eos:>
靴 を 履 い た 。 <:eos:>
靴 を 履 い た 。 <:eos:>
靴 を 履 く 。 <:eos:> <:pad:>
靴 を 履 い た 。 <:eos:>
駅 か ら 歩 い た 。 <:eos:>
駅 か ら 歩 い た 。 <:eos:>
駅 か ら 歩 く 。 <:eos:> <:pad:>
駅 か ら 歩 い た 。 <:eos:>
駅 か ら 家 ま 歩 い た 。 <:eos:> <:pad:>
駅 か ら 家 歩 い た 。 <:eos:> <:pad:> <:pad:>
駅 か ら 家 ま で 歩 い た 。 <:eos:>
駅 か ら 家 ま で 歩 く 。 <:eos:> <:pad:>

今度は、大体データに近い生成を行うようになりました。

これで「テキストファイルから学習データを読み込んで設定ファイルから作成したLLMで事前学習を行い、テストする」一連の流れを確認できたので、データの数や長さ・複雑さを変えながら試していきたいと思います。


創作プログラミングの街