test_tinystories.py

import torch
import math
import os
import yaml
import numpy as np
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tokenizers import Tokenizer
from util import LMDataset, get_token_id, generate_text, build_token_stream
from bnn import LuminaNet


config = "config/tinystories.yaml"
with open(config, "r") as f:
    cfg: dict = yaml.load(f, Loader=yaml.FullLoader)

TRAIN_TOKENS = "data/token_train.npy"
VAL_TOKENS = "data/token_val.npy"
TOKENIZER_PATH = "tokenizer/tokenizer.json"
SEQ_LEN = 32
BATCH_SIZE = cfg["Global"]["batch_size"]
device = torch.device(cfg["Global"]["device"] if torch.cuda.is_available() else "cpu")

print("[INFO] Loading tokenizer...")
tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
vocab_size = tokenizer.get_vocab_size()
PAD_ID = get_token_id("<unk>", tokenizer)
BOS_ID = get_token_id("<s>", tokenizer)
EOS_ID = get_token_id("</s>", tokenizer)
print(f"[INFO] Special ids -> PAD:{PAD_ID}, BOS:{BOS_ID}, EOS:{EOS_ID}")
print(f"[INFO] Vocab size = {vocab_size}")


def evaluate(model, loader):
    from tqdm import tqdm

    model.eval()

    total_loss = 0.0
    correct_top1 = 0
    correct_top3 = 0
    correct_top5 = 0
    total_samples = 0

    with torch.no_grad():
        for x, y in tqdm(loader):
            x, y = x.to(device), y.to(device)

            logits = model(x)
            next_token_logits = logits[:, -1, :]
            if y.dim() > 1:
                target = y[:, -1].long()
            else:
                target = y.long()

            loss = F.cross_entropy(next_token_logits, target, reduction="sum")
            total_loss += loss.item()

            _, topk_preds = next_token_logits.topk(5, dim=1)

            target_reshaped = target.view(-1, 1)

            correct_tensor = topk_preds.eq(target_reshaped)
            correct_top1 += correct_tensor[:, 0].sum().item()
            correct_top3 += correct_tensor[:, :3].sum().item()
            correct_top5 += correct_tensor.sum().item()

            total_samples += x.size(0)

    avg_loss = total_loss / total_samples

    try:
        perplexity = math.exp(avg_loss)
    except OverflowError:
        perplexity = float("inf")

    acc_top1 = correct_top1 / total_samples
    acc_top3 = correct_top3 / total_samples
    acc_top5 = correct_top5 / total_samples

    print(
        f"\n[Eval] Loss: {avg_loss:.4f} | PPL: {perplexity:.2f} | "
        f"Top-1 Acc: {acc_top1:.2%} | Top-3 Acc: {acc_top3:.2%} | Top-5 Acc: {acc_top5:.2%}"
    )

    return avg_loss, perplexity, acc_top1, acc_top5


def main():
    val_tokens = np.load(VAL_TOKENS)
    val_ds = LMDataset(val_tokens, SEQ_LEN)
    val_loader = DataLoader(
        val_ds, batch_size=BATCH_SIZE, shuffle=False, drop_last=True
    )

    name = cfg["Global"]["task_name"]
    net = LuminaNet(cfg)
    before = sum(p.numel() for p in net.parameters()) / 1000000
    net.load_model(ckpt_dir=f"ckpt/tinystories", filename=name)

    after = sum(p.numel() for p in net.parameters()) / 1000000
    print(f"Total parameters: {before:.2f}M -> {after:.2f}M")
    print("\n", generate_text(net, tokenizer, device, SEQ_LEN, EOS_ID), "\n")
    evaluate(net, val_loader)
    net.visualize_structure(name)


def build():
    if not os.path.exists(VAL_TOKENS):
        _ = build_token_stream("validation", VAL_TOKENS, tokenizer, EOS_ID)
    else:
        print(f"[INFO] Found existing {VAL_TOKENS}, skipping rebuild.")


if __name__ == "__main__":
    build()
    main()