| """Measure decision accuracy, NLL, Brier score and calibration error.""" |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| import tempfile |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
| from .data import ChoiceDataset, make_batch |
| from .atomic import atomic_output |
|
|
|
|
| class PredictionRecords: |
| def __init__(self, path, count): |
| self.path, self.count = Path(path), count |
|
|
| def __iter__(self): |
| with self.path.open(encoding="utf-8") as source: |
| for line in source: |
| yield json.loads(line) |
|
|
| def __len__(self): |
| return self.count |
|
|
|
|
| @torch.no_grad() |
| def collect(model, dataset, batch_size=8, cache=None, output_path=None): |
| model.eval() |
| if output_path is None: |
| directory=tempfile.TemporaryDirectory(prefix='paya-predictions-') |
| records=collect(model,dataset,batch_size,cache,output_path=Path(directory.name)/'predictions.jsonl') |
| records._temporary=directory |
| return records |
| if output_path is not None: |
| count = 0 |
| with atomic_output(output_path, "w") as output: |
| for rows in dataset.batches(batch_size): |
| result = model(make_batch(model, rows), cache=cache) |
| for row, logits in zip(rows, result["logits"]): |
| output.write(json.dumps({"logits": logits[:len(row["criteria"])].cpu().tolist(), |
| "target": row["target"], "criteria": row["criteria"], |
| **{k: row[k] for k in ("source_id", "group_id", "recording_id")}, |
| **{k: row[k] for k in ("mapping_id", "datasetname") if k in row}}, allow_nan=False)+"\n") |
| count += 1 |
| return PredictionRecords(output_path, count) |
|
|
|
|
| def metrics(records, temperature=1.0): |
| if not np.isfinite(temperature) or temperature <= 0: |
| raise ValueError("Need predictions and a finite positive temperature") |
| count = 0 |
| losses = briers = correct = 0.0 |
| bin_count = np.zeros(15, dtype=np.int64) |
| bin_correct = np.zeros(15) |
| bin_confidence = np.zeros(15) |
| for row in records: |
| z = np.asarray(row["logits"], dtype=np.float64)/temperature |
| if not np.isfinite(z).all(): |
| raise ValueError("Nonfinite prediction logits") |
| z -= z.max() |
| p = np.exp(z); p /= p.sum() |
| target = row["target"] |
| losses += float(np.log(np.exp(z).sum())-z[target]) |
| briers += float(((p - np.eye(len(p))[target])**2).sum()) |
| winner = min(range(len(p)), key=lambda i: (-p[i], row["criteria"][i])) |
| hit = float(winner == target) |
| confidence = float(p[winner]) |
| correct += hit |
| count += 1 |
| index = min(int(confidence*15),14) |
| bin_count[index] += 1 |
| bin_correct[index] += hit |
| bin_confidence[index] += confidence |
| if not count: |
| raise ValueError("Need predictions and a finite positive temperature") |
| reliability = [{"bin": i, "count": int(bin_count[i]), |
| "accuracy": float(bin_correct[i]/bin_count[i]), |
| "confidence": float(bin_confidence[i]/bin_count[i])} |
| for i in range(15) if bin_count[i]] |
| return {"count": count, "accuracy": correct/count, "nll": losses/count, |
| "brier": briers/count, "ece_15": float(sum(r["count"]/count*abs(r["accuracy"]-r["confidence"]) |
| for r in reliability)), |
| "reliability": reliability} |
|
|
|
|
| if __name__ == "__main__": |
| from .checkpoint import load_checkpoint |
| from .features import open_feature_cache |
| from .model import PayaModel |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", type=Path, default=Path("config/paya.yaml")) |
| parser.add_argument("--checkpoint", type=Path, required=True) |
| parser.add_argument("--split", choices=("selection", "calibration", "test"), default="test") |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--device", default="cpu") |
| args = parser.parse_args() |
| config = yaml.safe_load(args.config.read_text(encoding="utf-8")) |
| model = PayaModel.from_assets(config["assets"], config.get("preparation"), config.get("head"), args.device) |
| load_checkpoint(args.checkpoint, model) |
| dataset = ChoiceDataset(config["data"][args.split]) |
| cache = open_feature_cache(config) |
| records = collect(model, dataset, config.get("batch_size", 8), cache, |
| output_path=args.output.with_suffix(".predictions.jsonl")) |
| report = {"split": args.split, "calibrated": model.calibration is not None, |
| "metrics": metrics(records, float(model.temperature)), "predictions_path": str(records.path)} |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| with atomic_output(args.output,"w") as output: |
| json.dump(report,output,indent=2) |
| print(json.dumps(report["metrics"], indent=2)) |
|
|
