File size: 2,509 Bytes
87a0bf5 | 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 | """Save per-request outputs and decoding configuration after generation."""
import hashlib
import json
import os
from pathlib import Path
import re
def atomic(path, value):
path = Path(path)
temporary = path.with_suffix(path.suffix + '.tmp')
temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2))
temporary.replace(path)
def capture(ns):
root = Path(os.environ['ADSD_RUN_OUTPUT'])
index = ns['args'].sample_offset + ns['progress'] - 1
ids = ns['outputs'][0].detach().cpu().tolist()
prompt = ns['input_ids'][0].detach().cpu().tolist()
continuation = ids[len(prompt):]
text = ns['tokenizer1'].decode(continuation, skip_special_tokens=True)
def last_number(value):
matches = re.findall(r'-?\d+(?:\.\d+)?', value.replace(',', ''))
return matches[-1] if matches else None
counts = ns.get('counts', {}) if ns['args'].speculative else {}
row = dict(index=index, question=ns['q'], reference=ns['a'],
prompt_ids=prompt, generated_ids=continuation, generated_text=text,
elapsed_seconds=ns['end']-ns['start'], counts=counts,
correct_last_number=last_number(text) == last_number(ns['a']),
reached_cap=len(continuation) >= ns['args'].max_new_tokens,
prompt_sha256=hashlib.sha256(json.dumps(prompt).encode()).hexdigest())
atomic(root / f'row_{index:04d}.json', row)
atomic(root / 'progress.json', dict(status='running', completed=ns['progress'],
expected=ns['num_samples']))
def policy(ns):
root = Path(os.environ['ADSD_RUN_OUTPUT'])
import torch
import transformers
import signal
def interrupted(signum, _frame):
atomic(root / 'interrupted.json', {'signal': signum,
'completed_rows': len(list(root.glob('row_*.json'))),
'recovery': 'Retain partial rows; restart in a new output directory.'})
raise SystemExit(128 + signum)
signal.signal(signal.SIGTERM, interrupted)
signal.signal(signal.SIGINT, interrupted)
atomic(root / 'effective_policy.json', {
'torch': torch.__version__, 'transformers': transformers.__version__,
'models': {name: dict(dtype=str(ns[name].dtype),
attention=ns[name].config._attn_implementation,
generation_config=ns[name].generation_config.to_dict())
for name in ('model1', 'model2')},
'args': vars(ns['args'])})
|