Convert the editorial dataset into paragraph and sentence splits.
Remove images and headings, placing each heading's boundary on the preceding content.
For sentence inputs, place each paragraph's boundary on its final sentence.
Use the fixed 80/10/10 splits per outlet in configs/news_splits.json.
content_rows(rec)
Return content and heading rows, moving heading boundaries to the preceding content.
Source code in src/topic_segmentation/news_dataset/prepare_splits.py
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35 | def content_rows(rec):
"""Return content and heading rows, moving heading boundaries to the preceding content."""
block_types = rec.get("block_types", ["element"] * len(rec["body"]))
rows = [(text, label, heading)
for text, label, heading, kind in zip(rec["body"], rec["labels"], rec["title_labels"], block_types)
if kind != "image"]
if not rows:
return None
texts, labels, headings = map(list, zip(*rows))
for i, heading in enumerate(headings):
if heading == 1:
if i > 0 and headings[i - 1] == 0:
labels[i - 1] = 1
labels[i] = -100
if headings[-1] == 0 and labels[-1] != 1:
labels[-1] = 1
return list(zip(texts, labels, headings))
|
sentence_rows(rows)
Source code in src/topic_segmentation/news_dataset/prepare_splits.py
38
39
40
41
42
43
44
45
46 | def sentence_rows(rows):
out = []
for text, label, heading in rows:
if heading == 1:
out.append((text, -100, 1))
continue
parts = [s.strip() for s in sent_tokenize(text, language="english") if s.strip()]
out.extend((part, label if j == len(parts) - 1 else 0, 0) for j, part in enumerate(parts))
return out
|
build_record(rec, sentences)
Source code in src/topic_segmentation/news_dataset/prepare_splits.py
49
50
51
52
53
54
55
56
57
58
59
60
61 | def build_record(rec, sentences):
rows = content_rows(rec)
if rows and sentences:
rows = sentence_rows(rows)
if not rows:
return None
kept = [(text, 0 if label == -100 else label) for text, label, heading in rows if heading == 0]
return {
"articleId": rec["articleId"],
"outlet": rec["outlet"],
"sentences": [text for text, _ in kept],
"labels": [label for _, label in kept],
}
|
write_jsonl(path, examples)
Source code in src/topic_segmentation/news_dataset/prepare_splits.py
| def write_jsonl(path, examples):
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
for ex in examples:
f.write(json.dumps(ex, ensure_ascii=False) + "\n")
|
main()
Source code in src/topic_segmentation/news_dataset/prepare_splits.py
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88 | def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", default=str(ROOT / "data/dataset.jsonl"))
ap.add_argument("--out_root", default=str(ROOT / "prepared"))
args = ap.parse_args()
with open(args.src, "r", encoding="utf-8") as f:
raw = [json.loads(line) for line in f]
for name, sentences in VARIANTS:
records = {r["articleId"]: r for r in (build_record(rec, sentences) for rec in raw) if r is not None}
assert records.keys() == {article for ids in SPLITS.values() for article in ids}
for split, ids in SPLITS.items():
write_jsonl(os.path.join(args.out_root, name, f"{split}.jsonl"), [
{"articleId": article, "sentences": records[article]["sentences"], "labels": records[article]["labels"],
"title_labels": [0] * len(records[article]["sentences"]),
"section_topic_labels": [], "sentence_topic_labels": []}
for article in ids])
print(f"{name}: " + ", ".join(f"{split}={len(ids)}" for split, ids in SPLITS.items()))
|