diff --git a/scripts/fetch_corpus.sh b/scripts/fetch_corpus.sh new file mode 100755 index 0000000..0af0284 --- /dev/null +++ b/scripts/fetch_corpus.sh @@ -0,0 +1,4 @@ +#!/bin/bash +set -euo pipefail +mkdir -p corpus/raw +git clone --depth 1 https://github.com/suttacentral/bilara-data corpus/raw/bilara-data diff --git a/src/buddhagpt/corpus.py b/src/buddhagpt/corpus.py new file mode 100644 index 0000000..4c6e878 --- /dev/null +++ b/src/buddhagpt/corpus.py @@ -0,0 +1,31 @@ +import json, re +from pathlib import Path +from typing import Iterator + +def _seg_sort_key(k: str): + # "mn21:1.2" -> numeric-aware ordering of the segment path + # Handle mixed types (integers and strings) by wrapping in tuples + parts = [] + for p in re.split(r"[:.]", k): + if p.isdigit(): + parts.append((0, int(p))) # 0 for int, sorts before strings + else: + parts.append((1, p)) # 1 for string + return parts + +def segments_to_text(segments: dict[str, str]) -> str: + ordered = [segments[k] for k in sorted(segments, key=_seg_sort_key)] + return " ".join(s.strip() for s in ordered if s.strip()) + +def load_bilara_suttas(root: Path) -> Iterator[dict]: + base = root / "translation/en/sujato/sutta" + for f in sorted(base.rglob("*_translation-en-sujato.json")): + segments = json.loads(f.read_text()) + uid = f.name.split("_")[0] + title = next(iter(segments.values()), uid) + yield { + "uid": uid, + "title": title.strip(), + "text": segments_to_text(segments), + "collection": f.parent.name, + } diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_corpus.py b/tests/test_corpus.py new file mode 100644 index 0000000..c15f5c7 --- /dev/null +++ b/tests/test_corpus.py @@ -0,0 +1,17 @@ +import json +from pathlib import Path +from buddhagpt.corpus import segments_to_text, load_bilara_suttas + +def test_segments_to_text_joins_in_key_order(): + segs = {"mn21:1.2": "second.", "mn21:1.1": "First."} + assert segments_to_text(segs) == "First. second." + +def test_load_bilara_suttas(tmp_path): + d = tmp_path / "translation/en/sujato/sutta/mn" + d.mkdir(parents=True) + (d / "mn21_translation-en-sujato.json").write_text( + json.dumps({"mn21:0.1": "Middle Discourses 21", "mn21:1.1": "So I have heard."}) + ) + suttas = list(load_bilara_suttas(tmp_path)) + assert suttas[0]["uid"] == "mn21" + assert "So I have heard." in suttas[0]["text"]