| |
| """Score saved logits against the same official teacher-forced reference.""" |
| import argparse |
| import json |
| import numpy as np |
| import torch |
|
|
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("reference") |
| p.add_argument("--teacher-ids", required=True) |
| p.add_argument("--candidate", action="append", required=True, help="NAME:FILE.f32") |
| p.add_argument("--vocab", type=int, default=157184) |
| args = p.parse_args() |
| torch.set_num_threads(4) |
| teacher = [int(x) for x in open(args.teacher_ids).read().split()] |
| def read(path): |
| result = torch.from_numpy(np.fromfile(path, dtype=np.float32).reshape(len(teacher), args.vocab)).double() |
| if not torch.isfinite(result).all(): |
| raise ValueError(path) |
| return result |
| reference = read(args.reference) |
| log_p = reference.log_softmax(-1) |
| indices = torch.arange(len(teacher)) |
| summary = {"steps": len(teacher), "reference_nll": -log_p[indices, teacher].mean().item(), "candidates": {}} |
| for item in args.candidate: |
| name, path = item.split(":", 1) |
| actual = read(path) |
| log_q = actual.log_softmax(-1) |
| cosine = torch.nn.functional.cosine_similarity(reference, actual, dim=-1) |
| kl = (log_p.exp() * (log_p - log_q)).sum(-1) |
| summary["candidates"][name] = { |
| "nll": -log_q[indices, teacher].mean().item(), "kl_mean": kl.mean().item(), |
| "kl_max": kl.max().item(), "cosine_mean": cosine.mean().item(), |
| "cosine_min": cosine.min().item(), |
| "top1_agreement": int((actual.argmax(-1) == reference.argmax(-1)).sum()), |
| } |
| print(json.dumps(summary, indent=2)) |
|
|
