複数文と文末の学習
一つの文を1回学習してみたので、次はLLMで複数の文を学習/推論する仕組みを作ってみます。
複数のデータを扱うDatasetとTokenizer
LLMにデータを供給するDatasetでは、データサイズを返す__len__()とデータを返す__getitem__()という2つの関数を定義していました。
複数のデータを扱うには、データ数を__len__()で返し__getitem__()では引数で渡される番号(データ配列のインデックス)に応じたデータを返すようにします。
今回は、「私は歩く」「彼が走る」という二つの文を以下のDatasetで扱ってみます。
class TestDataset(Dataset):
def __init__(self, tokenizer):
self.data_list = []
texts = ["私は歩く", "彼が走る"]
for text in texts:
input_ids = tokenizer.encode(text, 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]
文を構成する語も増えので、トークナイザーも以下のような設定ファイルから生成するようにします。
{
"version": "1.0",
"truncation": null,
"padding": null,
"normalizer": null,
"pre_tokenizer": null,
"post_processor": null,
"model": {
"type": "Unigram",
"vocab": [
["<:pad:>", 0],
["私", 1],
["は", 2],
["歩", 3],
["く", 4],
["彼", 5],
["が", 6],
["走", 7],
["る", 8]
]
}
}
設定ファイルからトークナイザーを作ったら、モデルの設定ファイルも「vocab_size」を9に変更します。続けて、以下のコードで1回学習してみましょう。
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 = []
texts = ["私は歩く", "彼が走る"]
for text in texts:
input_ids = tokenizer.encode(text, 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=2, 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()
学習結果からの複数文の生成
学習が終了したら文を生成してみます。
今回は「私」「私は」「彼」「彼が」の4つの文に続く分をそれぞれ4つずつ生成してみることにしました。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("output/checkpoint-2")
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=6, pad_token_id=tokenizer.pad_token_id, do_sample=True, num_return_sequences=4)
for text in texts:
print(tokenizer.decode(text, skip_special_tokens = True))
私 歩 く は
私 歩 く は
私 は
私 は
私 は 歩 く く く く 歩
私 は は
私 は く く 歩 く く く
私 は は
彼 走 る 走 が が 走
彼 が 走 る る 走 走
彼 が る 走 が る 走
彼 走 走 走 走 が る
彼 が 走 く が 走 る 走
彼 が が 私 走 走 る 走
彼 が 走 る 走 る が る
彼 が 走 走 る る る 走
何となく繋がっている部分もある感じでしょうか。
ただ、だらだら生成が続くのがいまいちですね。文末に「。」と文の終わりを表すEOSトークンを追加してみましょう。
文末処理
まず、トークナイザーに「。」「<:eos:>」という2つのトークンを設定し、トークナイザー生成時の処理でもEOSを指定します。
{
"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]
]
}
}
この設定ファイルをtokenizer_src.jsonとして保存し以下のコードで処理すると、tokenizerディレクトリにEOSが設定されたトークナイザーが保存されます(トークナイザーのeos_token_idに<:eos:>のコード値が入る)。
from transformers import PreTrainedTokenizerFast
tokenizer = PreTrainedTokenizerFast(tokenizer_file="tokenizer_src.json", pad_token="<:pad:>", eos_token="<:eos:>")
tokenizer.save_pretrained(save_directory="tokenizer")
以下のように文末を変更したDatasetを作って学習してみましょう。
class TestDataset(Dataset):
def __init__(self, tokenizer):
self.data_list = []
texts = ["私は歩く。<:eos:>", "彼が走る。<:eos:>"]
for text in texts:
input_ids = tokenizer.encode(text, 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]
モデルの設定ファイルの「vocab_size」を11に設定してから学習コードを実行して1回だけ学習させます。続いて、以下のようにEOSなどのスペシャルトークンもデコードするようにしたテストコードでテストしてみます。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("output/checkpoint-2")
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=8, 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:> <:pad:>
私 歩 <:eos:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:>
私 は く 歩 は く 。 。 く
私 歩 く <:eos:> <:pad:> <:pad:> <:pad:> <:pad:> <:pad:>
私 は 。 歩 歩 く く く <:eos:> <:pad:>
私 は 歩 く は く 。 く 。 。
私 は 。 歩 く く <:eos:> <:pad:> <:pad:> <:pad:>
私 は 歩 く 。 歩 <:eos:> <:pad:> <:pad:> <:pad:>
彼 が る 走 走 る く 走 私
彼 走 が 走 走 走 る る る
彼 走 <:pad:> る が 走 る 走 <:eos:>
彼 が 走 る が る る <:eos:> <:pad:>
彼 が 走 走 る 。 が る <:eos:>
彼 が る 走 。 走 走 走 <:eos:>
彼 が が る 走 走 。 走 <:eos:>
彼 が 走 る る 走 <:eos:> <:pad:> <:pad:>
いまいちですね。
学習コードのTrainingArgumentsを以下のように変更して2回学習させてみましょうか。
training_args = TrainingArguments(output_dir="output", save_steps=4, num_train_epochs=2, learning_rate=5e-5, lr_scheduler_type="constant", per_device_train_batch_size=1)
テストプログラムを「output/checkpoint-4」からモデルを読み込むように変更して試してみると…
私 は 。 く 。 。 。 く <:eos:>
私 は 。 く 。 。 は 。 <:eos:>
私 く は 。 。 。 く 歩 。
私 は 。 。 く 。 <:eos:> <:pad:> <:pad:>
私 は 。 。 が 。 歩 。 。 く
私 は 。 く は く 。 。 。 。
私 は 歩 く 。 。 。 。 。 。
私 は 。 。 く が 。 。 。 。
彼 が が が 走 走 走 。 。
彼 が が 。 。 <:eos:> <:pad:> <:pad:> <:pad:>
彼 走 が 。 走 る 。 <:eos:> <:pad:>
彼 が 走 走 。 走 。 る <:eos:>
彼 が が 走 走 。 。 <:eos:> <:pad:> <:pad:>
彼 が が る 。 。 。 走 。 る
彼 が が 走 走 。 走 。 。 走
彼 が 走 る る 。 。 。 。 <:eos:>
あまり変わってないような…
学習を4回(2つの文を4回なので計8文)にしてみましょう。
私 は 歩 歩 。 <:eos:>
私 は 歩 歩 。 <:eos:>
私 は 歩 歩 。 <:eos:>
私 は 歩 歩 歩 <:eos:>
私 は 歩 歩 歩 <:eos:> <:pad:>
私 は 歩 。 <:eos:> <:pad:> <:pad:>
私 は 歩 歩 歩 。 <:eos:>
私 は 歩 歩 歩 <:eos:> <:pad:>
彼 が 走 る 。 <:eos:> <:pad:>
彼 が 走 走 る 。 <:eos:>
彼 が 走 る 。 <:eos:> <:pad:>
彼 が 走 る 。 <:eos:> <:pad:>
彼 が 走 る 。 <:eos:> <:pad:> <:pad:>
彼 が 走 る 。 <:eos:> <:pad:> <:pad:>
彼 が 走 走 る 。 走 <:eos:>
彼 が 走 る る 。 <:eos:> <:pad:>
「彼が走る。」の方は、大体うまく行くようになりました。8回学習すると…
私 は 歩 く 。 <:eos:>
私 は 歩 く 。 <:eos:>
私 は 歩 く 。 <:eos:>
私 は 歩 く 。 <:eos:>
私 は 歩 く 。 <:eos:>
私 は 歩 く 。 <:eos:>
私 は 歩 く 。 <:eos:>
私 は 歩 く 。 <:eos:>
彼 が 走 る 。 <:eos:>
彼 が 走 る 。 <:eos:>
彼 が 走 る 。 <:eos:>
彼 が 走 る 。 <:eos:>
彼 が 走 る 。 <:eos:>
彼 が 走 る 。 <:eos:>
彼 が 走 る 。 <:eos:>
彼 が 走 る 。 <:eos:>
学習させた文そのままを出力するようになりました。
学習時に末尾にEOSを付けないと、8回学習させてもこんな感じになります。
私 は 歩 く 。 。 。 。 。
私 は 歩 く 。 。 。 。 。
私 は 歩 く 。 。 。 。 。
私 は 歩 く 。 。 。 。 。
私 は 歩 く 。 。 。 。 。 。
私 は 歩 く 。 。 。 。 。 。
私 は 歩 く 。 。 。 。 。 。
私 は 歩 く 。 。 。 。 。 。
彼 が 走 る 。 。 。 。 。
彼 が 走 る 。 。 。 。 。
彼 が 走 る 。 。 。 。 。
彼 が 走 る 。 。 。 。 。
彼 が 走 る 。 。 。 。 。 。
彼 が 走 る 。 。 。 。 。 。
彼 が 走 る 。 。 。 。 。 。
彼 が 走 る 。 。 。 。 。 。