Photo by
featurecompass
サクッとフルスクラッチでLLM作成してみる！
今回は、サクッとフルスクラッチでLLMの作成をする記事を書いてみます。
環境は、Google ColabのTPU v2-8で実行となります。
1.必要なライブラリのインストール
!pip install torch numpy tqdm2.データの準備
import torch
import numpy as np
from torch.utils.data import DataLoader, Dataset
import random
# シェイクスピアのデータを使用（トイデータセット）
text = """To be, or not to be, that is the question:
Whether 'tis nobler in the mind to suffer
The slings and arrows of outrageous fortune,
Or to take arms against a sea of troubles"""
# 文字ベースのトークン化（文字レベルの言語モデルとして構築）
chars = sorted(list(set(text)))
vocab_size = len(chars)
# 文字を整数に変換する辞書とその逆
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
# トークン化関数
def tokenize(text):
return [stoi[ch] for ch in text]
# テキストを整数シーケンスに変換
data = tokenize(text)
# シーケンスデータをバッチに変換（修正版）
def get_batch(data, block_size, batch_size):
ix = random.randint(0, len(data) - block_size - 1)
inputs = []
targets = []
for i in range(batch_size):
# シーケンスの終わりを超える場合は、ランダムな位置から取得
if ix + block_size + 1 > len(data):
ix = random.randint(0, len(data) - block_size - 1)
inputs.append(data[ix:ix+block_size])
targets.append(data[ix+1:ix+block_size+1])
ix += block_size
return torch.tensor(inputs), torch.tensor(targets)
block_size = 8 # 各トレーニングシーケンスの長さ
batch_size = 4 # バッチサイズ3.Transformerベースのモデルを定義
import torch.nn as nn
import torch.nn.functional as F
class SimpleTransformer(nn.Module):
def __init__(self, vocab_size, block_size, embedding_dim):
super(SimpleTransformer, self).__init__()
self.embed = nn.Embedding(vocab_size, embedding_dim)
self.transformer = nn.Transformer(d_model=embedding_dim, nhead=4, num_encoder_layers=2)
self.fc_out = nn.Linear(embedding_dim, vocab_size)
def forward(self, x):
# x: (batch_size, block_size)
x = self.embed(x) # (batch_size, block_size, embedding_dim)
x = x.permute(1, 0, 2) # Transformer expects (block_size, batch_size, embedding_dim)
x = self.transformer(x, x)
x = x.permute(1, 0, 2) # Back to (batch_size, block_size, embedding_dim)
logits = self.fc_out(x) # (batch_size, block_size, vocab_size)
return logits
# ハイパーパラメータ
embedding_dim = 128
model = SimpleTransformer(vocab_size, block_size, embedding_dim)
# モデルを確認
print(model)
4.トレーニング
# オプティマイザと損失関数を設定
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()
# トレーニング
epochs = 5000 # エポック数
for epoch in range(epochs):
model.train()
# バッチデータを取得
inputs, targets = get_batch(data, block_size, batch_size)
# 順伝播
logits = model(inputs)
loss = loss_fn(logits.view(-1, vocab_size), targets.view(-1))
# 逆伝播とオプティマイザの更新
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 100 == 0:
print(f"Epoch {epoch} Loss: {loss.item()}")
print("トレーニング完了")5.テキスト生成
# テキスト生成関数
def generate(model, start_text, max_new_tokens):
model.eval()
start_tokens = tokenize(start_text)
start_tokens = torch.tensor(start_tokens).unsqueeze(0) # (1, len(start_tokens))
generated = start_tokens.tolist()[0]
for _ in range(max_new_tokens):
logits = model(start_tokens)
probs = F.softmax(logits[:, -1, :], dim=-1)
next_token = torch.multinomial(probs, num_samples=1).item()
generated.append(next_token)
start_tokens = torch.tensor(generated).unsqueeze(0)
return ''.join([itos[token] for token in generated])
# プロンプトを設定してテキスト生成
start_text = "To be"
generated_text = generate(model, start_text, max_new_tokens=100)
print(f"Generated text:\n{generated_text}")
実行結果は、次のようになります。
Generated text:
To beb lbo rr oob e e lrbuer orbro boseeooe o olr r l or reunnr efrser ehello,ren i biwrbl e leel 説明
データセット準備: テキストをトークン化し、文字レベルのデータに変換しています。このモデルは文字ベースでシーケンス生成を行うため、非常にシンプルです。
Transformerモデル: シンプルなTransformerベースのモデルを構築し、埋め込み層、トランスフォーマーレイヤー、出力層を定義しています。
トレーニング: トレーニングループでモデルを学習し、損失を最小化します。クロスエントロピー損失を使って、予測と実際の文字列との誤差を計算しています。
テキスト生成: トレーニング後、プロンプトに基づいて新しいテキストを生成します。