File size: 4,627 Bytes
e2c9d42 | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | #!/usr/bin/env python3
"""Build a split summary that preserves RAGTruth's own train/test boundary.
RAGTruth ships a train/test split. That boundary is authoritative and is kept
exactly: every row RAGTruth calls `test` becomes our held-out eval set, and
nothing else does. Reshuffling it would make our numbers incomparable to every
published RAGTruth result, including LLM-Check's.
Our protocol additionally needs a validation split, which RAGTruth does not
provide, for checkpoint and threshold selection. It is carved out of RAGTruth's
*train* portion only, grouped by `source_id`, so the eval set is untouched and
no source document appears on both sides of the train/validation line. Several
responses share one article; splitting by row would leak an article's content
between fitting and selection.
Usage:
python scripts/make_ragtruth_split.py \
--input ~/Desktop/datasets/hallucinations/ragtruth_llama2_7b_chat.jsonl \
--output ~/Desktop/datasets/hallucinations/ragtruth_llama2_7b_chat_split.json \
--validation-fraction 0.15 --seed 42
"""
from __future__ import annotations
import argparse
import collections
import json
from pathlib import Path
import sys
import numpy as np
REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from scripts.make_eval_split import (
NATIVE_SPLIT_SCHEME,
build_native_split,
dataset_binding,
)
def parse_args(argv=None):
p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--input", type=Path, required=True)
p.add_argument("--output", type=Path, required=True)
p.add_argument(
"--validation-fraction",
type=float,
default=0.15,
help="Share of RAGTruth's TRAIN groups held out for validation.",
)
p.add_argument("--seed", type=int, default=42)
return p.parse_args(argv)
def main(argv=None):
args = parse_args(argv)
rows = [json.loads(l) for l in args.input.open(encoding="utf-8")]
if not rows:
raise SystemExit("empty input")
# Built by the SAME function validate_split_summary() uses to verify it,
# so the writer and the checker cannot drift apart.
built = build_native_split(
rows,
seed=args.seed,
validation_fraction=args.validation_fraction,
native_split_field="metadata.split",
eval_value="test",
grouping_key="prompt",
)["split"]
train_indices = built["train_indices"]
validation_indices = built["validation_indices"]
eval_indices = built["eval_indices"]
native_test = {i for i, r in enumerate(rows) if r["metadata"]["split"] == "test"}
assert set(eval_indices) == native_test, "eval must equal RAGTruth test exactly"
train_sources = {rows[i]["metadata"]["extra"]["ragtruth_source_id"] for i in train_indices}
val_sources = {rows[i]["metadata"]["extra"]["ragtruth_source_id"] for i in validation_indices}
assert not (train_sources & val_sources), "source_id leaked across train/validation"
summary = {
"scheme": NATIVE_SPLIT_SCHEME,
"seed": args.seed,
"n": len(rows),
"dataset_binding": dataset_binding(rows),
"test_fraction": len(eval_indices) / len(rows),
"validation_fraction": args.validation_fraction,
"native": {
"split_field": "metadata.split",
"eval_value": "test",
"grouping_key": "prompt",
},
"provenance": {
"eval_split": "RAGTruth's own `test` rows, preserved exactly",
"validation_split": "carved from RAGTruth's `train` rows only, grouped by source_id",
},
"split": {
"train_indices": train_indices,
"validation_indices": validation_indices,
"eval_indices": eval_indices,
},
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
def h_rate(idx):
tok = sum(len(rows[i]["labels"]) for i in idx)
h = sum(sum(1 for v in rows[i]["labels"] if v == 1) for i in idx)
return tok, h / max(tok, 1)
print(f"wrote {args.output}")
for name, idx in (("train", train_indices), ("validation", validation_indices), ("eval(test)", eval_indices)):
tok, rate = h_rate(idx)
print(f" {name:<12} {len(idx):>5} rows {tok:>8,} tok H {rate:.2%}")
print(f" source groups: train {len(train_sources)}, validation {len(val_sources)}, disjoint")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|