File size: 4,702 Bytes
3d7c238 | 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """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())
|