Add 3x-context dataset variant (trainset --radius)

- build_trainset gains --radius (chars each side of the cited name) and --out;
  merge-gap scales with radius. Default 600 unchanged.
- trainset_3x + train/val/test_3x.jsonl: same 10,519 triples and same trust split,
  but ~3x more surrounding prose per triple (~47 -> ~132 tokens/triple, median
  ~3.7k tokens/sample). Keeps the 100% name-in-text guarantee.
- DATASET.md documents both context sizes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Florian Herzog 2026-06-10 16:37:30 +02:00
parent e35d98c2cd
commit 9dc870b8d0
6 changed files with 708 additions and 12 deletions

View File

@ -13,6 +13,12 @@ target that cannot be derived from its text.
## Files (this is all you need — no rebuild required)
Two context sizes are provided, identical triples and splits, differing only in
how much surrounding prose is kept around each cited entity (the `_3x` variant has
~3× more distractor text per triple — useful to test robustness to longer input).
**Standard context (~600 chars each side, ~47 tokens/triple):**
| File | Samples | Size | Purpose |
|------|--------:|-----:|---------|
| `data/rdf_poc/trainset.jsonl` | 332 | ~6 MB | full set (all samples) |
@ -20,6 +26,18 @@ target that cannot be derived from its text.
| `data/rdf_poc/val.jsonl` | 35 | ~0.8 MB | validation split |
| `data/rdf_poc/test.jsonl` | 35 | ~0.5 MB | test split |
**3× context (~1800 chars each side, ~132 tokens/triple, median ~3.7k tokens/sample):**
| File | Samples | Size |
|------|--------:|-----:|
| `data/rdf_poc/trainset_3x.jsonl` | 332 | ~9 MB |
| `data/rdf_poc/train_3x.jsonl` | 262 | ~7 MB |
| `data/rdf_poc/val_3x.jsonl` | 35 | ~1 MB |
| `data/rdf_poc/test_3x.jsonl` | 35 | ~0.8 MB |
Both variants keep the 100 % name-in-text guarantee. The `_3x` set is produced by
`build_rdf_dataset.py trainset --radius 1800 --out ...`; any other radius works too.
Splits are **by trust (CIK)**: all funds of one trust stay in one split, so the
model cannot memorise a trust's providers from another split (no leakage). The
assignment is a deterministic hash of the CIK and is reproducible.

View File

@ -1201,16 +1201,19 @@ def _name_positions(label, low):
return pos
def _evidence_excerpt(text_norm_full, raw_text, label, predicate=None):
def _evidence_excerpt(text_norm_full, raw_text, label, predicate=None, radius=None):
"""Cut an excerpt of raw_text around a REAL occurrence of the entity name.
Among the positions where the name genuinely appears (see _name_positions),
prefer one with a role keyword (custodian/adviser/transfer agent/...) nearby
-- the actual service-provider statement -- but ALWAYS fall back to the first
real occurrence so the name is guaranteed to be inside the excerpt (the
training input must contain every target). Returns (excerpt, start) or
(None, -1) only when the name does not occur at all (a real grounding error).
training input must contain every target). `radius` chars are taken on each
side (default EXCERPT_RADIUS); a larger radius yields more surrounding context
(distractor text) per triple. Returns (excerpt, start) or (None, -1) only when
the name does not occur at all (a real grounding error).
"""
rad = radius if radius is not None else EXCERPT_RADIUS
low = raw_text.lower()
occ = _name_positions(label, low)
if not occ:
@ -1221,15 +1224,19 @@ def _evidence_excerpt(text_norm_full, raw_text, label, predicate=None):
if any(k in win for k in _ROLE_HINT):
pos = p
break
a = max(0, pos - EXCERPT_RADIUS)
b = min(len(raw_text), pos + EXCERPT_RADIUS)
a = max(0, pos - rad)
b = min(len(raw_text), pos + rad)
return raw_text[a:b].strip(), a
def build_trainset():
def build_trainset(radius=None, out_path=None):
"""Assemble the training-ready dataset: per sample, the relevant TEXT EXCERPTS
combined with ONLY the grounded (extractable) triples as the target.
`radius` (chars each side of the name) controls how much surrounding context
is kept per triple; default EXCERPT_RADIUS. A 3x radius produces a variant with
~3x more text per triple (more distractor context before/after the cited span).
For each labelled sample (samples_labelled.jsonl):
- keep only triples with extractable=True (every target is in the text);
- for each, cut a coherent excerpt of the prospectus around the entity name
@ -1243,9 +1250,11 @@ def build_trainset():
log.error("run `label` first (need %s)", SAMPLES_LABELLED_PATH); return
gold = {json.loads(l)["cik"]: json.loads(l)
for l in open(GOLD_PATH, encoding="utf-8")}
rad = radius if radius is not None else EXCERPT_RADIUS
dst = Path(out_path) if out_path else TRAINSET_PATH
n = n_skip = n_drop = 0
tot_in = tot_tr = 0
with open(TRAINSET_PATH, "w", encoding="utf-8") as out:
with open(dst, "w", encoding="utf-8") as out:
for line in open(SAMPLES_LABELLED_PATH, encoding="utf-8"):
r = json.loads(line)
ents = gold.get(r["cik"], {}).get("entities", {})
@ -1263,7 +1272,7 @@ def build_trainset():
kept = []
for t in cand:
lbl = ents.get(t["o"], {}).get("label", t["o"])
exc, a = _evidence_excerpt(None, full, lbl)
exc, a = _evidence_excerpt(None, full, lbl, radius=rad)
if exc is None:
n_drop += 1 # name not in text -> not a valid sample
continue
@ -1274,8 +1283,9 @@ def build_trainset():
continue
spans.sort()
merged = []
gap = max(200, rad // 3) # merge nearby windows; scale gap with radius
for a, b in spans:
if merged and a <= merged[-1][1] + 200:
if merged and a <= merged[-1][1] + gap:
merged[-1] = (merged[-1][0], max(merged[-1][1], b))
else:
merged.append((a, b))
@ -1299,7 +1309,7 @@ def build_trainset():
}
out.write(json.dumps(rec, ensure_ascii=False) + "\n")
n += 1; tot_in += len(excerpt); tot_tr += len(kept)
print(f"\nTrainset: {n} samples -> {TRAINSET_PATH} ({n_skip} NONE-samples skipped)")
print(f"\nTrainset: {n} samples -> {dst} (radius={rad}, {n_skip} NONE-samples skipped)")
if n:
print(f" grounded triples (targets): {tot_tr:,}")
print(f" dropped (LLM-grounded but name absent from prose): {n_drop} "
@ -1330,7 +1340,11 @@ def main():
"(default, enforces text-extractability); name=name anywhere; "
"none=keep all, only annotate")
sub.add_parser("label") # label extractability (needs prior match run)
sub.add_parser("trainset") # assemble training set: excerpts + grounded triples
ts = sub.add_parser("trainset") # assemble training set: excerpts + grounded triples
ts.add_argument("--radius", type=int, default=None,
help="chars of context each side of the name (default 600)")
ts.add_argument("--out", default=None,
help="output path (default data/rdf_poc/trainset.jsonl)")
sp = sub.add_parser("split")
sp.add_argument("--val-frac", type=float, default=0.10)
sp.add_argument("--test-frac", type=float, default=0.10)
@ -1357,7 +1371,7 @@ def main():
elif args.cmd == "label":
build_label()
elif args.cmd == "trainset":
build_trainset()
build_trainset(radius=args.radius, out_path=args.out)
elif args.cmd == "split":
build_split(val_frac=args.val_frac, test_frac=args.test_frac, src=args.src)
elif args.cmd == "all":

File diff suppressed because one or more lines are too long

262
data/rdf_poc/train_3x.jsonl Normal file

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

35
data/rdf_poc/val_3x.jsonl Normal file

File diff suppressed because one or more lines are too long