用几万个英语句子构建一个自己的小语言模型

pip install tokenizers datasets

data/sentences.txt

I am learning English.
This is a sample sentence.
Language models are powerful.
...

from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace

# 初始化 tokenizer
tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
tokenizer.pre_tokenizer = Whitespace()

# BPE 训练器配置
trainer = BpeTrainer(
    vocab_size=3000,          # 可根据语料调整大小
    show_progress=True,
    special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
)

# 训练 tokenizer
files = ["data/sentences.txt"]
tokenizer.train(files, trainer)

# 保存 tokenizer 到本地
tokenizer.save("bpe_tokenizer.json")

from tokenizers import Tokenizer

# 加载 tokenizer
tokenizer = Tokenizer.from_file("bpe_tokenizer.json")

# 编码单句
output = tokenizer.encode("I am learning English.")
print("Token IDs:", output.ids)
print("Tokens:", output.tokens)

# 假设你要构建一个 language model 的训练集，输入 = 前面 token，输出 = 下一个 token

def build_dataset(tokenizer, input_file, block_size=64):
    input_ids = []
    with open(input_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            ids = tokenizer.encode(line).ids
            input_ids.extend(ids + [tokenizer.token_to_id("[SEP]")])  # 作为句子结束标志

    # 切分成多个 block
    blocks = []
    for i in range(0, len(input_ids) - block_size, block_size):
        x = input_ids[i:i+block_size]
        y = input_ids[i+1:i+block_size+1]
        blocks.append((x, y))
    return blocks

import torch
import torch.nn as nn

class TinyTransformer(nn.Module):
    def __init__(self, vocab_size, emb_size=128, block_size=64):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, emb_size)
        self.position_embedding = nn.Embedding(block_size, emb_size)
        self.ln = nn.LayerNorm(emb_size)
        self.lm_head = nn.Linear(emb_size, vocab_size)

    def forward(self, idx):
        B, T = idx.shape
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))
        x = tok_emb + pos_emb
        x = self.ln(x)
        logits = self.lm_head(x)
        return logits

model = TinyTransformer(vocab_size=tokenizer.get_vocab_size())
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(10):
    for x, y in blocks:  # blocks 为训练样本
        x = torch.tensor([x])
        y = torch.tensor([y])
        logits = model(x)
        loss = loss_fn(logits.view(-1, logits.size(-1)), y.view(-1))
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    print(f"Epoch {epoch}: Loss = {loss.item()}")

def generate(model, tokenizer, start_text="I am", max_len=20):
    model.eval()
    tokens = tokenizer.encode(start_text).ids
    input_ids = torch.tensor([tokens])
    for _ in range(max_len):
        logits = model(input_ids)
        next_token = torch.argmax(logits[:, -1, :], dim=-1).item()
        input_ids = torch.cat([input_ids, torch.tensor([[next_token]])], dim=1)
    print("Generated:", tokenizer.decode(input_ids[0].tolist()))