複数文と文末の学習

一つの文を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回学習させてもこんな感じになります。

私 は 歩 く 。 。 。 。 。
私 は 歩 く 。 。 。 。 。
私 は 歩 く 。 。 。 。 。
私 は 歩 く 。 。 。 。 。
私 は 歩 く 。 。 。 。 。 。
私 は 歩 く 。 。 。 。 。 。
私 は 歩 く 。 。 。 。 。 。
私 は 歩 く 。 。 。 。 。 。
彼 が 走 る 。 。 。 。 。
彼 が 走 る 。 。 。 。 。
彼 が 走 る 。 。 。 。 。
彼 が 走 る 。 。 。 。 。
彼 が 走 る 。 。 。 。 。 。
彼 が 走 る 。 。 。 。 。 。
彼 が 走 る 。 。 。 。 。 。
彼 が 走 る 。 。 。 。 。 。

創作プログラミングの街