Spaces:
Running
Running
| """Build the request stream from TAB (public train split, MIT): one request = one sentence of an ECHR judgment. | |
| Documents are shuffled (seed 0) and split by document: HOLDOUT_DOCS documents give the spot-check holdout, the rest give | |
| the stream. Gold spans (TAB annotator 1, the three format types) are kept for reference only; the demo learns from the | |
| teacher, never from gold. | |
| usage: prep.py TAB.parquet OUT_DIR | |
| """ | |
| import json | |
| import random | |
| import re | |
| import sys | |
| import pyarrow.parquet as pq | |
| TYPES = {"DATETIME", "CODE", "QUANTITY"} | |
| N_STREAM, N_HOLDOUT, HOLDOUT_DOCS = 3000, 400, 120 | |
| ABBR = r"(?<!\bMr)(?<!\bMrs)(?<!\bMs)(?<!\bDr)(?<!\bno)(?<!\bNo)(?<!\bnos)(?<!\bNos)(?<!\bpp)(?<!\bp)(?<!\bpara)(?<!\bparas)(?<!\bArt)(?<!\bv)(?<!\bSt)(?<!\bJr)(?<!\bi\.e)(?<!\be\.g)(?<!\bseq)(?<!\bop)(?<!\bcit)(?<!\b[A-Z])" | |
| SENT_END = re.compile(ABBR + r"[.;](?=\s+[A-Z(\"'])") | |
| def sentences(text): | |
| out = [] | |
| for m in re.finditer(r"[^\n]+", text): | |
| start, para = m.start(), m.group() | |
| s = 0 | |
| for e in SENT_END.finditer(para): | |
| out.append((start + s, start + e.end())) | |
| s = e.end() | |
| while s < len(para) and para[s] == " ": | |
| s += 1 | |
| if s < len(para): | |
| out.append((start + s, start + len(para))) | |
| return out | |
| def mentions(ann): | |
| for v in (ann or {}).values(): | |
| if isinstance(v, dict) and v.get("entity_mentions"): | |
| return v["entity_mentions"] | |
| return [] | |
| def main(pqf, out): | |
| docs = pq.read_table(pqf).to_pylist() | |
| random.Random(0).shuffle(docs) | |
| rows = {"holdout": [], "stream": []} | |
| for k, d in enumerate(docs): | |
| part = "holdout" if k < HOLDOUT_DOCS else "stream" | |
| ann = d["annotations"] | |
| if isinstance(ann, str): | |
| ann = json.loads(ann) | |
| ents = [(m["start_offset"], m["end_offset"], m["entity_type"]) for m in mentions(ann) if m.get("entity_type") in TYPES] | |
| for a, b in sentences(d["text"]): | |
| t = d["text"][a:b].strip() | |
| if not (40 <= len(t) <= 400) or t.isupper(): | |
| continue | |
| a2 = d["text"].index(t, a) | |
| gold = sorted({(s - a2, e - a2, ty) for s, e, ty in ents if s >= a2 and e <= a2 + len(t)}) | |
| rows[part].append({"doc_id": d["doc_id"], "text": t, | |
| "gold": [{"start": s, "end": e, "type": ty, "text": t[s:e]} for s, e, ty in gold]}) | |
| rng = random.Random(1) | |
| sel = {"stream": rng.sample(rows["stream"], N_STREAM), "holdout": rng.sample(rows["holdout"], N_HOLDOUT)} | |
| with open(f"{out}/requests.jsonl", "w") as f: | |
| for part, rs in sel.items(): | |
| for i, r in enumerate(rs): | |
| f.write(json.dumps({"id": f"{part[0]}{i:04d}", "part": part, **r}) + "\n") | |
| for part, rs in sel.items(): | |
| print(part, len(rs), "of", len(rows[part]), "docs", len({r["doc_id"] for r in rs}), | |
| "with gold", sum(1 for r in rs if r["gold"]), "gold spans", sum(len(r["gold"]) for r in rs)) | |
| if __name__ == "__main__": | |
| main(sys.argv[1], sys.argv[2]) | |