86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
"""Resumable evaluation: checkpoint plugin.
|
|
|
|
Run-level checkpointing: every completed sample is appended to a jsonl
|
|
checkpoint file; on restart (crash/OOM/service down), completed samples are
|
|
restored and ONLY the missing ones are generated. The report is then
|
|
assembled from restored + fresh predictions.
|
|
|
|
from evalharness.eval.checkpoint import CheckpointStore
|
|
store = CheckpointStore(path='run.jsonl', model=model_spec, dataset=name)
|
|
done = store.load() # {sample_key: prediction-dict}
|
|
... generate only missing ...
|
|
store.append(sample_key, pred) # after EACH sample (crash-safe)
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional
|
|
|
|
|
|
class CheckpointStore:
|
|
"""Append-only jsonl checkpoint; keyed by stable sample identity."""
|
|
|
|
def __init__(self, path: str, model: str = '', dataset: str = ''):
|
|
self.path = Path(path)
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.model = model
|
|
self.dataset = dataset
|
|
self._entries: Dict[str, Dict[str, Any]] = {}
|
|
self._fh = None
|
|
|
|
@staticmethod
|
|
def key_for(sample, idx: int) -> str:
|
|
"""Stable per-sample key: prefer dataset-native ids, fall back to
|
|
a hash of the input text (survives re-orderings)."""
|
|
native = (sample.metadata or {}).get('task_id') \
|
|
or (sample.metadata or {}).get('id') \
|
|
or (sample.metadata or {}).get('instance_id')
|
|
if native:
|
|
return str(native)
|
|
import hashlib
|
|
|
|
h = hashlib.md5((sample.input_text or '').encode('utf-8')).hexdigest()[:16]
|
|
return f'{idx}:{h}'
|
|
|
|
def load(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Read all checkpointed predictions (idempotent)."""
|
|
self._entries = {}
|
|
if not self.path.exists():
|
|
return self._entries
|
|
with open(self.path, encoding='utf-8') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
rec = json.loads(line)
|
|
self._entries[rec['key']] = rec.get('pred', {})
|
|
except (ValueError, KeyError):
|
|
continue # torn tail line from a crash -- safe to skip
|
|
return self._entries
|
|
|
|
def append(self, key: str, pred: Dict[str, Any]) -> None:
|
|
"""Persist one prediction immediately (fsync-free append is fine:
|
|
worst case loses the last in-flight sample on crash)."""
|
|
rec = {'key': key, 'ts': time.time(), 'pred': pred}
|
|
with open(self.path, 'a', encoding='utf-8') as f:
|
|
f.write(json.dumps(rec, ensure_ascii=False) + '\n')
|
|
self._entries[key] = pred
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._entries)
|
|
|
|
def summary(self) -> Dict[str, Any]:
|
|
return {'path': str(self.path), 'restored': len(self._entries),
|
|
'size_kb': self.path.stat().st_size // 1024 if self.path.exists() else 0}
|
|
|
|
|
|
def checkpoint_path(root: str, dataset: str, model: str, tag: str = '') -> str:
|
|
"""Deterministic path per (dataset, model[, tag]) so re-runs resume."""
|
|
import hashlib
|
|
|
|
h = hashlib.md5(f'{dataset}|{model}|{tag}'.encode()).hexdigest()[:10]
|
|
return os.path.join(root, 'ckpt', f'{dataset}-{h}.jsonl')
|