File size: 2,904 Bytes
d9eec2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """Phase 3 — évaluation : perplexité + génération d'échantillons.
python -m src.evaluate --ckpt checkpoints/text_50m/best.pt \\
--data-dir data/processed/wikipedia_fr --eval-batches 50 --sample
"""
from __future__ import annotations
import argparse
import json
import math
from pathlib import Path
import numpy as np
import torch
from src.data.prepare_text import load_bin
from src.generate import generate_text
from src.model.text_model import CortexTextLM, TextConfig
from src.tokenizer import load_tokenizer
from src.train.train_text import TrainConfig, estimate_loss, pick_autocast, pick_device
def load_checkpoint(path: str | Path, device: torch.device) -> tuple[CortexTextLM, dict]:
ckpt = torch.load(path, map_location=device, weights_only=False)
model = CortexTextLM(TextConfig(**ckpt["text_config"])).to(device)
model.load_state_dict(ckpt["model"])
model.eval()
return model, ckpt
@torch.no_grad()
def perplexity(model: CortexTextLM, data: np.memmap, cfg: TrainConfig,
device: torch.device, autocast_ctx) -> float:
metrics = estimate_loss(model, {"val": data}, cfg, device, autocast_ctx)
return math.exp(min(metrics["val"], 20)), metrics["val"]
def main() -> None:
parser = argparse.ArgumentParser(description="Évaluer un checkpoint texte Cortex-0.")
parser.add_argument("--ckpt", type=Path, required=True)
parser.add_argument("--data-dir", type=Path, required=True)
parser.add_argument("--split", default="val")
parser.add_argument("--batch-size", type=int, default=16)
parser.add_argument("--eval-batches", type=int, default=50)
parser.add_argument("--device", default="auto")
parser.add_argument("--sample", action="store_true")
parser.add_argument("--prompt", default="La France est")
parser.add_argument("--max-new-tokens", type=int, default=60)
args = parser.parse_args()
device = pick_device(args.device)
model, ckpt = load_checkpoint(args.ckpt, device)
data = load_bin(Path(args.data_dir) / f"{args.split}.bin")
cfg = TrainConfig(batch_size=args.batch_size, eval_batches=args.eval_batches,
ctx_len=model.config.ctx_len, device=args.device)
autocast_ctx, _ = pick_autocast(device, "float32")
ppl, loss = perplexity(model, data, cfg, device, autocast_ctx)
result = {"checkpoint": str(args.ckpt), "step": ckpt.get("step"),
"split": args.split, "val_loss": loss, "perplexity": ppl,
"n_params": model.count_parameters()}
if args.sample:
tokenizer = load_tokenizer()
result["samples"] = [
generate_text(model, tokenizer, args.prompt, max_new_tokens=args.max_new_tokens,
temperature=t, top_k=40)
for t in (0.7, 1.0)
]
print(json.dumps(result, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()
|