Correct the model card (pooling is at the LABELS suffix's last token); add code/classifier from project commit 4ced7554f
3d7c238 verified | """Predict with a trained run. Run as ``python -m classifier.predict`` (pod, or any rented GPU for serving). | |
| Input: an export split (``--split dev|benchmark``) or any JSONL of export-format rows (``--input``, e.g. the | |
| pool or the robustness set). Output rows: ``{"id", "split", "n_tokens", "truncated", "logits"}`` with raw logits | |
| per field in taxonomy option order; ``classifier.calibrate`` turns them into calibrated, gated probabilities. | |
| The summary line reports posts/s and tokens/s (the serving throughput) and, with ``--repeat N``, the largest | |
| probability difference between repeats (the determinism check). | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| import time | |
| from .calibrate import softmax | |
| from .data import inference_batches, load_jsonl, load_split, write_jsonl | |
| from .formats import encode | |
| from .taxonomy import Taxonomy | |
| def main(argv=None): | |
| p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) | |
| p.add_argument('--run', required=True, help='Run directory written by classifier.train') | |
| p.add_argument('--which', default='best', choices=['best', 'last']) | |
| source = p.add_mutually_exclusive_group(required=True) | |
| source.add_argument('--split', choices=['train', 'dev', 'benchmark']) | |
| source.add_argument('--input', help='JSONL of export-format rows') | |
| p.add_argument('--data', help='Export directory (default: the run config\'s)') | |
| p.add_argument('--model', help='Checkpoint directory (default: the run config\'s)') | |
| p.add_argument('--out', required=True) | |
| p.add_argument('--token-budget', type=int, default=65536, help='Padded tokens per batch') | |
| p.add_argument('--limit', type=int, default=0) | |
| p.add_argument('--repeat', type=int, default=1) | |
| p.add_argument('--device', default='cuda') | |
| p.add_argument('--kernels', default='fast', choices=['fast', 'torch']) | |
| args = p.parse_args(argv) | |
| import torch | |
| from .model import build, load_tokenizer, predict, tokenizer_fn | |
| run = Path(args.run) | |
| cfg = json.loads((run / 'config.json').read_text()) | |
| if args.model: | |
| cfg['model'] = args.model | |
| data = args.data or cfg['data'] | |
| tax = Taxonomy.load(run / 'taxonomy.json') | |
| device = torch.device(args.device) | |
| ckpt = run / args.which | |
| model = build(cfg, tax.sizes(), device, train=False, adapter_dir=ckpt / 'adapter', heads_path=ckpt / 'heads.pt') | |
| from .check_env import fast_path, torch_path | |
| if args.kernels == 'torch': | |
| torch_path().__enter__() | |
| kernels = fast_path(model) | |
| if args.kernels == 'fast' and not kernels['ok']: | |
| raise SystemExit('Fast kernels are not active: %s' % json.dumps(kernels)) | |
| tok = load_tokenizer(cfg['model']) | |
| rows = load_split(data, args.split) if args.split else load_jsonl(args.input) | |
| if args.limit: | |
| rows = rows[:args.limit] | |
| t = time.time() | |
| encoded = [encode(r, tokenizer_fn(tok), cfg['max_len']) for r in rows] | |
| ids = [e[0] for e in encoded] | |
| encode_seconds = time.time() - t | |
| batches = inference_batches([len(x) for x in ids], args.token_budget) | |
| runs = [] | |
| for _ in range(max(1, args.repeat)): | |
| t = time.time() | |
| logits = predict(model, ids, tok.pad_token_id, batches, device, tax.names) | |
| if device.type == 'cuda': | |
| torch.cuda.synchronize() | |
| runs.append((logits, time.time() - t)) | |
| logits, seconds = runs[0] | |
| out_rows = [{'id': r['id'], 'split': r.get('split'), 'n_tokens': encoded[i][1]['n_tokens'], | |
| 'truncated': encoded[i][1]['cut'], 'logits': logits[i]} for i, r in enumerate(rows)] | |
| write_jsonl(args.out, out_rows) | |
| tokens = sum(len(x) for x in ids) | |
| summary = dict(rows=len(rows), seconds=round(seconds, 2), encode_seconds=round(encode_seconds, 2), | |
| posts_per_second=round(len(rows) / seconds, 1), tokens_per_second=round(tokens / seconds), | |
| tokens=tokens, truncated_rows=sum(1 for e in encoded if e[1]['cut']), which=args.which, | |
| checkpoint=str(ckpt), kernels=args.kernels) | |
| if len(runs) > 1: | |
| diff, same = 0.0, True | |
| for other, _ in runs[1:]: | |
| for i in logits: | |
| for f in tax.names: | |
| a, b = softmax(logits[i][f]), softmax(other[i][f]) | |
| diff = max(diff, max(abs(x - y) for x, y in zip(a, b))) | |
| same &= a.index(max(a)) == b.index(max(b)) | |
| summary.update(max_prob_diff=diff, argmax_identical=same) | |
| if device.type == 'cuda': | |
| summary['max_memory_gb'] = round(torch.cuda.max_memory_allocated() / 2 ** 30, 1) | |
| print(json.dumps(summary)) | |
| return 0 | |
| if __name__ == '__main__': | |
| raise SystemExit(main()) | |