Photo by
featurecompass
ゼロから作る超簡易的なLLM-Google Colab編
Google Colabでゼロから作る超簡易的なLLM作成の仕方を紹介します。
Google Colabを開き、ランタイムはT4 GPUとします。
次のコードを貼り付けます。
####################################
# 1. 必要なライブラリのインストールとインポート
####################################
# !pip install torch
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)
####################################
# 2. 学習用のテキストデータの用意（サンプル）
####################################
# 簡単のための短い文字列。実際には自由に増やしたり、大きなコーパスを使ってください。
text_data = """\
こんにちは。これは簡単な言語モデルのテストです。
同じような文章を繰り返して学習し、どの程度再現できるかを確認します。
こんにちは。これは2回目の文章です。
"""
# ここでは文字レベルでトークナイズします
chars = sorted(list(set(text_data)))
vocab_size = len(chars)
# 文字 -> インデックス、インデックス -> 文字 の対応表を作る
char2idx = {ch: i for i, ch in enumerate(chars)}
idx2char = {i: ch for i, ch in enumerate(chars)}
print("Vocabulary size:", vocab_size)
print("Vocabulary chars:", chars)
# テキストをインデックス列に変換
encoded_data = [char2idx[c] for c in text_data]
####################################
# 3. データセットの定義
####################################
# 入力シーケンス長 seq_len の次に来る文字を予測するようにする。
# 例）"こん" -> 'に'を予測
class TextDataset(Dataset):
def __init__(self, data, seq_len=20):
self.data = data
self.seq_len = seq_len
def __len__(self):
return len(self.data) - self.seq_len
def __getitem__(self, idx):
# シーケンス部分
x = self.data[idx:idx+self.seq_len]
# 次の1文字(トークン)
y = self.data[idx+1:idx+self.seq_len+1]
x = torch.tensor(x, dtype=torch.long)
y = torch.tensor(y, dtype=torch.long)
return x, y
seq_len = 20
dataset = TextDataset(encoded_data, seq_len=seq_len)
batch_size = 16
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)
####################################
# 4. モデルの定義 (Transformerベースの簡易LLM)
####################################
class SimpleTransformerModel(nn.Module):
def __init__(self, vocab_size, embed_dim, num_heads, hidden_dim, num_layers, seq_len):
super().__init__()
self.seq_len = seq_len
# 文字をベクトルに埋め込む
self.embedding = nn.Embedding(vocab_size, embed_dim)
# Positional Encoding (非常に簡易的に実装)
self.pos_embedding = nn.Embedding(seq_len, embed_dim)
# トランスフォーマーブロックを積む
encoder_layer = nn.TransformerEncoderLayer(
d_model=embed_dim,
nhead=num_heads,
dim_feedforward=hidden_dim,
batch_first=True
)
self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
# 最後に出力層で次の文字確率を出す
self.fc_out = nn.Linear(embed_dim, vocab_size)
def forward(self, x):
# x: [batch_size, seq_len]
batch_size, seq_len = x.shape
# 埋め込み: [batch_size, seq_len, embed_dim]
x_emb = self.embedding(x)
# 位置情報の埋め込み
positions = torch.arange(0, seq_len, device=x.device).unsqueeze(0) # [1, seq_len]
pos_emb = self.pos_embedding(positions) # [1, seq_len, embed_dim]
x_emb = x_emb + pos_emb # 位置情報を足し合わせる
# Transformer Encoder
transformed = self.transformer_encoder(x_emb) # [batch_size, seq_len, embed_dim]
# 出力層
logits = self.fc_out(transformed) # [batch_size, seq_len, vocab_size]
return logits
# ハイパーパラメータ（本当に簡易的に）
embed_dim = 64
num_heads = 4
hidden_dim = 128
num_layers = 2
model = SimpleTransformerModel(
vocab_size=vocab_size,
embed_dim=embed_dim,
num_heads=num_heads,
hidden_dim=hidden_dim,
num_layers=num_layers,
seq_len=seq_len
).to(device)
####################################
# 5. 学習の設定
####################################
criterion = nn.CrossEntropyLoss() # 次の文字を当てる分類問題
optimizer = optim.Adam(model.parameters(), lr=1e-3)
epochs = 10 # 回数を増やせばより学習するが、Colabの時間と相談して調整
####################################
# 6. 学習ループ
####################################
for epoch in range(epochs):
total_loss = 0
for x_batch, y_batch in dataloader:
x_batch = x_batch.to(device)
y_batch = y_batch.to(device)
optimizer.zero_grad()
# logits: [batch_size, seq_len, vocab_size]
logits = model(x_batch)
# 予測対象は各時刻の文字なので、(batch_size*seq_len) x vocab_size と (batch_size*seq_len) 形式に潰す
loss = criterion(
logits.view(-1, vocab_size),
y_batch.view(-1)
)
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(dataloader)
print(f"Epoch [{epoch+1}/{epochs}] Loss: {avg_loss:.4f}")
####################################
# 7. テキスト生成用の関数
####################################
def generate_text(model, start_text="こんにちは。", max_length=100):
model.eval()
# start_text を文字列からインデックス列に変換
input_indices = [char2idx.get(ch, 0) for ch in start_text]
# すでに start_text 分の長さがあるので、そこから続ける
generated = input_indices[:] # 今までの出力
# 生成を開始
for _ in range(max_length):
# 直近の seq_len 分をモデルに入れる (足りなければ前を0埋め)
x = generated[-seq_len:] if len(generated) > seq_len else generated
# バッチ化
x_tensor = torch.tensor([x], dtype=torch.long, device=device)
# 必要に応じて先頭に padding を入れる (簡単のため先頭を切り落とす実装にしている)
if x_tensor.size(1) < seq_len:
# seq_len に満たない分を0埋め (実際はpad_tokenのIDにすべきだが簡易的に0に)
pad_size = seq_len - x_tensor.size(1)
x_tensor = torch.cat([
torch.zeros((1, pad_size), dtype=torch.long, device=device),
x_tensor
], dim=1)
with torch.no_grad():
# [1, seq_len, vocab_size]
logits = model(x_tensor)
# 最終トークンに対応する出力を取り出す
last_logit = logits[:, -1, :] # [1, vocab_size]
# 確率分布に変換 (sampling したい場合は softmax 後にサンプリング)
probs = torch.softmax(last_logit, dim=-1)
# 最大値をとるインデックスを次の文字とする (greedy)
next_idx = torch.argmax(probs, dim=-1).item()
generated.append(next_idx)
# インデックスから文字列に変換
generated_text = "".join(idx2char[i] for i in generated)
return generated_text
####################################
# 8. テキスト生成を試す
####################################
generated_sample = generate_text(model, start_text="こんにちは。", max_length=50)
print("Generated text sample:")
print(generated_sample)Google Colabで上記コードを実行しますと、次の結果が得られます。
こんにちは。 同じようです。 同じような文章を繰り返して学習して学習して学習して学習して学習して学習して学習して
今回は、学習回数のepochsが10でしたので、epochs=1000にして、上記を実行すると、ほぼ類似の文章の出力となります。
こんにちは。
こんにちは。これは。これは2回目の文章です。
同じような文章を繰り返して学習し、どの程度再現できるこのコードをベースに、学習用のテキストデータを長くしたり、epchsを多くしたりすると色々な結果が出力されて面白いので試してみるといいかもしれません。