From 9dd254f702cd7eea659dbee0ccbd259ba4ac76ac Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 1 Aug 2026 16:34:15 +0800 Subject: [PATCH 1/2] feat(data): real Tashkeela++ fetcher via HuggingFace Hub MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the env-var placeholder in scripts/fetch_data.sh with a proper Python fetcher that knows the canonical dataset locations. Primary source: Misraj/Sadeed_Tashkeela — gated, requires HF_TOKEN. Fallback: community-datasets/tashkeela — GPLv2 open access. The fetcher streams parquet → one-line-per-chunk text, skips blank and overlong lines, writes atomically via .tmp rename, and exits with a clear error message on gated-repo failures (including the URL the user must visit to grant access). Also adds pythonpath = ["src", "."] to pytest config so tests run without pip install -e ., and pyarrow>=15.0 to [publish] extras. Smoke-tested end-to-end: fetch 100 lines → RababaArabicData consumes them → (bare, diacritized) pairs ready for StudentTrainer. --- pyproject.toml | 2 + scripts/fetch_data.py | 237 +++++++++++++++++++++++++++++++++++++++ tests/test_fetch_data.py | 54 +++++++++ 3 files changed, 293 insertions(+) create mode 100755 scripts/fetch_data.py create mode 100644 tests/test_fetch_data.py diff --git a/pyproject.toml b/pyproject.toml index 3d1ac34..2dadd81 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ export = [ publish = [ "huggingface_hub>=0.23", "omegaconf>=2.3", + "pyarrow>=15.0", ] dev = [ "pytest>=8.0", @@ -52,6 +53,7 @@ include = ["framework*", "tasks*"] [tool.pytest.ini_options] testpaths = ["tests"] +pythonpath = ["src", "."] addopts = "-ra -q --strict-markers" markers = [ "slow: marks tests as slow (deselect with '-m \"not slow\"')", diff --git a/scripts/fetch_data.py b/scripts/fetch_data.py new file mode 100755 index 0000000..f4d791c --- /dev/null +++ b/scripts/fetch_data.py @@ -0,0 +1,237 @@ +"""Fetch raw corpora from HuggingFace Hub. + +Replaces ``scripts/fetch_data.sh``. The shell version required manual +env-var URLs; this one knows the canonical dataset locations and +validates downloads. + +For ``rababa_arabic``: + Default source: ``Misraj/Sadeed_Tashkeela`` — gated, requires HF_TOKEN + and access grant at + https://huggingface.co/datasets/Misraj/Sadeed_Tashkeela + Fallback: ``community-datasets/tashkeela`` — GPLv2, open access, raw + book text that needs heavier cleaning (handled by the data module). + +For ``rababa_hebrew`` and ``secryst_thai_ipa`` the upstream sources are +not yet on HF as datasets — leave the manual env-var path intact in +``fetch_data.sh`` until they are. +""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT / "src")) + +DEFAULT_OUT = ROOT / "data" / "raw" + +DATASETS = { + "rababa_arabic": { + "primary": { + "repo_id": "Misraj/Sadeed_Tashkeela", + "repo_type": "dataset", + "files": [ + "data/train-00000-of-00003.parquet", + "data/train-00001-of-00003.parquet", + "data/train-00002-of-00003.parquet", + ], + "test_files": ["data/test-00000-of-00001.parquet"], + "text_column": "text", + "out_name": "tashkeela_plus_plus.txt", + "split_lines": False, + "note": ( + "Gated dataset. Visit " + "https://huggingface.co/datasets/Misraj/Sadeed_Tashkeela, " + "log in, accept the terms, then export HF_TOKEN." + ), + }, + "fallback": { + "repo_id": "community-datasets/tashkeela", + "repo_type": "dataset", + "files": None, + "text_column": "text", + "out_name": "tashkeela_plus_plus.txt", + "split_lines": True, + "note": ( + "Open-access raw corpus (GPLv2). Each row is a full book; " + "we split on newlines and skip lines >1024 chars." + ), + }, + }, +} + + +def fetch_task( + task: str, + out_dir: Path, + max_samples: int | None, + use_fallback: bool, +) -> Path: + cfg = DATASETS.get(task) + if cfg is None: + raise SystemExit( + f"No fetcher registered for task '{task}'. " + f"Known: {sorted(DATASETS)}" + ) + source = cfg["fallback"] if use_fallback else cfg["primary"] + + try: + import importlib.util + + if importlib.util.find_spec("huggingface_hub") is None: + raise ImportError("huggingface_hub not installed") + except ImportError as e: + raise SystemExit( + "huggingface_hub is required. Install with: " + "pip install -e '.[publish]'" + ) from e + + token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / source["out_name"] + text_column = source["text_column"] + + files = source["files"] + if files is None: + files = _list_repo_files(source["repo_id"], source["repo_type"], token) + + count = _stream_parquet_to_text( + files=files, + repo_id=source["repo_id"], + repo_type=source["repo_type"], + token=token, + text_column=text_column, + out_path=out_path, + max_samples=max_samples, + split_lines=source.get("split_lines", False), + ) + print( + f"[{task}] wrote {count:,} lines -> {out_path} " + f"({out_path.stat().st_size:,} bytes) " + f"from {source['repo_id']}" + ) + if count == 0: + raise SystemExit( + f"No lines written. Source note: {source.get('note', '')}" + ) + return out_path + + +def _list_repo_files(repo_id: str, repo_type: str, token: str | None) -> list[str]: + from huggingface_hub import HfApi + + api = HfApi(token=token) + files = api.list_repo_files(repo_id, repo_type=repo_type) + return [f for f in files if f.endswith((".parquet", ".json", ".jsonl", ".txt"))] + + +def _stream_parquet_to_text( + files: list[str], + repo_id: str, + repo_type: str, + token: str | None, + text_column: str, + out_path: Path, + max_samples: int | None, + split_lines: bool = False, + max_line_chars: int = 1024, +) -> int: + """Stream the ``text_column`` of each parquet file to ``out_path``. + + Each row's text is written as one line (after whitespace collapse). + If ``split_lines`` is set (raw book corpora), the row's text is split + on embedded newlines first — one row may carry many verse-sized lines. + Lines longer than ``max_line_chars`` are skipped (training chunks + should be ~50-60 words; longer ones are typically misplits). + """ + import pyarrow.parquet as pq + from huggingface_hub import hf_hub_download + + written = 0 + tmp = out_path.with_suffix(".txt.tmp") + with tmp.open("w", encoding="utf-8") as fp: + for fpath in files: + local = hf_hub_download( + repo_id=repo_id, + filename=fpath, + repo_type=repo_type, + token=token, + ) + pf = pq.ParquetFile(local) + for batch in pf.iter_batches(batch_size=1024, columns=[text_column]): + col = batch.column(text_column).to_pylist() + for blob in col: + if not blob: + continue + chunks = blob.splitlines() if split_lines else [blob] + for raw in chunks: + line = " ".join(raw.split()) + if not line or len(line) > max_line_chars: + continue + fp.write(line) + fp.write("\n") + written += 1 + if max_samples is not None and written >= max_samples: + tmp.replace(out_path) + return written + tmp.replace(out_path) + return written + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--task", + required=True, + choices=sorted(DATASETS), + help="Which task corpus to fetch.", + ) + parser.add_argument( + "--out-dir", + type=Path, + default=DEFAULT_OUT, + help=f"Output directory (default: {DEFAULT_OUT}).", + ) + parser.add_argument( + "--max-samples", + type=int, + default=None, + help="Cap number of lines written (dev/CI mode).", + ) + parser.add_argument( + "--fallback", + action="store_true", + help=( + "Use the open-access fallback dataset instead of the primary " + "(useful when the primary is gated and no HF_TOKEN is set)." + ), + ) + args = parser.parse_args() + + try: + fetch_task( + task=args.task, + out_dir=args.out_dir, + max_samples=args.max_samples, + use_fallback=args.fallback, + ) + except SystemExit: + raise + except Exception as e: + msg = str(e) + if "GatedRepoError" in type(e).__name__ or "gated" in msg.lower(): + cfg = DATASETS[args.task] + note = cfg["primary"].get("note", "") + raise SystemExit( + f"Gated dataset: {type(e).__name__}\n{note}\n" + f"Or rerun with --fallback to use {cfg['fallback']['repo_id']}." + ) from e + raise SystemExit(f"Fetch failed: {type(e).__name__}: {msg}") from e + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_fetch_data.py b/tests/test_fetch_data.py new file mode 100644 index 0000000..cc76e1c --- /dev/null +++ b/tests/test_fetch_data.py @@ -0,0 +1,54 @@ +"""Tests for ``scripts/fetch_data.py``. + +Network access is NOT required — these tests cover import, CLI +argument handling, and the gated-repo error path. Real download is +covered by the smoke test in the README. +""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +SCRIPT = Path(__file__).resolve().parent.parent / "scripts" / "fetch_data.py" + + +@pytest.fixture(scope="module") +def fetch_module(): + spec = importlib.util.spec_from_file_location("fetch_data", SCRIPT) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_fetcher_imports(fetch_module) -> None: + assert hasattr(fetch_module, "fetch_task") + assert hasattr(fetch_module, "DATASETS") + assert "rababa_arabic" in fetch_module.DATASETS + + +def test_fetcher_primary_has_parquet_files(fetch_module) -> None: + primary = fetch_module.DATASETS["rababa_arabic"]["primary"] + assert primary["repo_id"] == "Misraj/Sadeed_Tashkeela" + assert primary["repo_type"] == "dataset" + assert len(primary["files"]) == 3 + assert all(f.endswith(".parquet") for f in primary["files"]) + + +def test_fetcher_fallback_uses_open_dataset(fetch_module) -> None: + fallback = fetch_module.DATASETS["rababa_arabic"]["fallback"] + assert fallback["repo_id"] == "community-datasets/tashkeela" + assert fallback["split_lines"] is True + + +def test_fetcher_unknown_task_errors(fetch_module, tmp_path: Path) -> None: + with pytest.raises(SystemExit, match="No fetcher registered"): + fetch_module.fetch_task( + task="not_a_task", + out_dir=tmp_path, + max_samples=None, + use_fallback=False, + ) From c2fc1870cc4a9b8be18c7f9c8ad3d56b2f9602af Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 1 Aug 2026 17:35:41 +0800 Subject: [PATCH 2/2] =?UTF-8?q?feat(data):=20switch=20to=20open=20arbml/ta?= =?UTF-8?q?shkeelav2=20=E2=80=94=20no=20gating,=20paired=20columns?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drops Misraj/Sadeed_Tashkeela (gated, requires manual approval) as the primary source. Replacement: arbml/tashkeelav2 — open-access, pre-split into train/test, and pre-paired (each row has both `text` and `diacratized` columns). No HF_TOKEN, no acceptance click-through. Output format is now TSV with two columns: `barediacritized`. The data module reads both directly — no in-pipeline stripping needed, and the dataset's canonical letter forms are preserved. Falls back to the legacy single-text path (.txt + strip_diacritics) when only raw community-datasets/tashkeela is available. Layout: primary = arbml/tashkeelav2 (TSV, open) fallback = community-datasets/tashkeela (TXT, GPLv2 raw) New tests: - fetcher primary is open + uses TSV - fetcher fallback still uses raw text path - data module reads TSV pairs (bare+diacritized) - data module falls back to legacy .txt when no TSV - data module prefers TSV when both exist 46 tests pass; smoke-tested against arbml/tashkeelav2 end-to-end. --- scripts/fetch_data.py | 186 ++++++++++++++++++++++---------- src/tasks/rababa_arabic/data.py | 35 +++++- tests/test_fetch_data.py | 21 ++-- tests/test_rababa_arabic.py | 62 +++++++++++ 4 files changed, 235 insertions(+), 69 deletions(-) diff --git a/scripts/fetch_data.py b/scripts/fetch_data.py index f4d791c..6909bc8 100755 --- a/scripts/fetch_data.py +++ b/scripts/fetch_data.py @@ -5,15 +5,26 @@ validates downloads. For ``rababa_arabic``: - Default source: ``Misraj/Sadeed_Tashkeela`` — gated, requires HF_TOKEN - and access grant at - https://huggingface.co/datasets/Misraj/Sadeed_Tashkeela - Fallback: ``community-datasets/tashkeela`` — GPLv2, open access, raw - book text that needs heavier cleaning (handled by the data module). + Default source: ``arbml/tashkeelav2`` — open-access, pre-split, and + pre-paired (each row carries both the bare and diacritized text). + No HF_TOKEN, no gating, no acceptance click-through. + Fallback: ``community-datasets/tashkeela`` — GPLv2, raw book text + that needs heavier cleaning (handled by the data module). For ``rababa_hebrew`` and ``secryst_thai_ipa`` the upstream sources are not yet on HF as datasets — leave the manual env-var path intact in ``fetch_data.sh`` until they are. + +Output format +------------- +The fetcher writes one of two file shapes: + + * ``.tsv`` — paired columns ``barediacritized``. Used + when the upstream dataset provides both columns (arbml/tashkeelav2). + The data module reads both directly — no stripping needed. + * ``.txt`` — one diacritized line per row. Used when the + upstream is unpaired (raw tashkeela). The data module strips + diacritics itself. """ from __future__ import annotations @@ -31,33 +42,34 @@ DATASETS = { "rababa_arabic": { "primary": { - "repo_id": "Misraj/Sadeed_Tashkeela", + "repo_id": "arbml/tashkeelav2", "repo_type": "dataset", - "files": [ - "data/train-00000-of-00003.parquet", - "data/train-00001-of-00003.parquet", - "data/train-00002-of-00003.parquet", + "train_files": [ + "data/train-00000-of-00002-74e12fe61707b796.parquet", + "data/train-00001-of-00002-4f28b1f4b7fd8dd8.parquet", ], - "test_files": ["data/test-00000-of-00001.parquet"], - "text_column": "text", - "out_name": "tashkeela_plus_plus.txt", - "split_lines": False, + "test_files": ["data/test-00000-of-00001-ffd46a42fa4bfebf.parquet"], + "bare_column": "text", + "diacritized_column": "diacratized", + "out_name": "tashkeela_plus_plus.tsv", "note": ( - "Gated dataset. Visit " - "https://huggingface.co/datasets/Misraj/Sadeed_Tashkeela, " - "log in, accept the terms, then export HF_TOKEN." + "Open-access, pre-paired, pre-split. ~116k train + 58k " + "test rows. Each row provides both bare + diacritized " + "text — no in-pipeline stripping required." ), }, "fallback": { "repo_id": "community-datasets/tashkeela", "repo_type": "dataset", - "files": None, - "text_column": "text", + "train_files": None, + "test_files": [], + "bare_column": None, + "diacritized_column": "text", "out_name": "tashkeela_plus_plus.txt", - "split_lines": True, "note": ( - "Open-access raw corpus (GPLv2). Each row is a full book; " - "we split on newlines and skip lines >1024 chars." + "Open-access raw corpus (GPLv2). Each row is a full " + "book; we split on newlines and skip lines >1024 chars. " + "Bare form is derived by the data module via strip()." ), }, }, @@ -78,36 +90,45 @@ def fetch_task( ) source = cfg["fallback"] if use_fallback else cfg["primary"] - try: - import importlib.util + import importlib.util - if importlib.util.find_spec("huggingface_hub") is None: - raise ImportError("huggingface_hub not installed") - except ImportError as e: + if importlib.util.find_spec("huggingface_hub") is None: raise SystemExit( "huggingface_hub is required. Install with: " "pip install -e '.[publish]'" - ) from e + ) token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") out_dir.mkdir(parents=True, exist_ok=True) out_path = out_dir / source["out_name"] - text_column = source["text_column"] - - files = source["files"] - if files is None: - files = _list_repo_files(source["repo_id"], source["repo_type"], token) - - count = _stream_parquet_to_text( - files=files, - repo_id=source["repo_id"], - repo_type=source["repo_type"], - token=token, - text_column=text_column, - out_path=out_path, - max_samples=max_samples, - split_lines=source.get("split_lines", False), - ) + + train_files = source["train_files"] + if train_files is None: + train_files = _list_repo_files(source["repo_id"], source["repo_type"], token) + + if source["bare_column"] and source["diacritized_column"]: + count = _stream_paired_tsv( + files=train_files, + repo_id=source["repo_id"], + repo_type=source["repo_type"], + token=token, + bare_column=source["bare_column"], + diacritized_column=source["diacritized_column"], + out_path=out_path, + max_samples=max_samples, + ) + else: + count = _stream_single_text( + files=train_files, + repo_id=source["repo_id"], + repo_type=source["repo_type"], + token=token, + text_column=source["diacritized_column"], + out_path=out_path, + max_samples=max_samples, + split_lines=True, + ) + print( f"[{task}] wrote {count:,} lines -> {out_path} " f"({out_path.stat().st_size:,} bytes) " @@ -128,7 +149,61 @@ def _list_repo_files(repo_id: str, repo_type: str, token: str | None) -> list[st return [f for f in files if f.endswith((".parquet", ".json", ".jsonl", ".txt"))] -def _stream_parquet_to_text( +def _stream_paired_tsv( + files: list[str], + repo_id: str, + repo_type: str, + token: str | None, + bare_column: str, + diacritized_column: str, + out_path: Path, + max_samples: int | None, + max_line_chars: int = 1024, +) -> int: + """Stream both columns to a TSV: ``barediacritized`` per line.""" + import pyarrow.parquet as pq + from huggingface_hub import hf_hub_download + + written = 0 + tmp = out_path.with_suffix(out_path.suffix + ".tmp") + cols = [bare_column, diacritized_column] + with tmp.open("w", encoding="utf-8") as fp: + for fpath in files: + local = hf_hub_download( + repo_id=repo_id, + filename=fpath, + repo_type=repo_type, + token=token, + ) + pf = pq.ParquetFile(local) + for batch in pf.iter_batches(batch_size=1024, columns=cols): + bare_list = batch.column(bare_column).to_pylist() + dia_list = batch.column(diacritized_column).to_pylist() + for bare, dia in zip(bare_list, dia_list, strict=True): + if not bare or not dia: + continue + bare = " ".join(bare.split()) + dia = " ".join(dia.split()) + if ( + not bare + or not dia + or len(bare) > max_line_chars + or len(dia) > max_line_chars + ): + continue + fp.write(bare) + fp.write("\t") + fp.write(dia) + fp.write("\n") + written += 1 + if max_samples is not None and written >= max_samples: + tmp.replace(out_path) + return written + tmp.replace(out_path) + return written + + +def _stream_single_text( files: list[str], repo_id: str, repo_type: str, @@ -139,19 +214,13 @@ def _stream_parquet_to_text( split_lines: bool = False, max_line_chars: int = 1024, ) -> int: - """Stream the ``text_column`` of each parquet file to ``out_path``. - - Each row's text is written as one line (after whitespace collapse). - If ``split_lines`` is set (raw book corpora), the row's text is split - on embedded newlines first — one row may carry many verse-sized lines. - Lines longer than ``max_line_chars`` are skipped (training chunks - should be ~50-60 words; longer ones are typically misplits). - """ + """Stream one text column to ``out_path`` (one line per row, or per + embedded line if ``split_lines`` is set for raw book corpora).""" import pyarrow.parquet as pq from huggingface_hub import hf_hub_download written = 0 - tmp = out_path.with_suffix(".txt.tmp") + tmp = out_path.with_suffix(out_path.suffix + ".tmp") with tmp.open("w", encoding="utf-8") as fp: for fpath in files: local = hf_hub_download( @@ -206,7 +275,8 @@ def main() -> int: action="store_true", help=( "Use the open-access fallback dataset instead of the primary " - "(useful when the primary is gated and no HF_TOKEN is set)." + "(useful when the primary is unavailable or you want GPLv2 " + "raw text)." ), ) args = parser.parse_args() @@ -223,11 +293,9 @@ def main() -> int: except Exception as e: msg = str(e) if "GatedRepoError" in type(e).__name__ or "gated" in msg.lower(): - cfg = DATASETS[args.task] - note = cfg["primary"].get("note", "") raise SystemExit( - f"Gated dataset: {type(e).__name__}\n{note}\n" - f"Or rerun with --fallback to use {cfg['fallback']['repo_id']}." + f"Gated dataset encountered: {type(e).__name__}\n" + f"Switch to --fallback to use an open dataset." ) from e raise SystemExit(f"Fetch failed: {type(e).__name__}: {msg}") from e return 0 diff --git a/src/tasks/rababa_arabic/data.py b/src/tasks/rababa_arabic/data.py index 4da5cce..9364846 100644 --- a/src/tasks/rababa_arabic/data.py +++ b/src/tasks/rababa_arabic/data.py @@ -73,8 +73,12 @@ class RababaArabicData(DataModule): def prepare_data(self) -> PreparedData: if self._prepared is not None: return self._prepared - raw_path = self.data_root / "raw" / f"{self.config.source}.txt" - examples = self._read_examples(raw_path) + tsv_path = self.data_root / "raw" / f"{self.config.source}.tsv" + txt_path = self.data_root / "raw" / f"{self.config.source}.txt" + if tsv_path.is_file(): + examples = self._read_examples_tsv(tsv_path) + else: + examples = self._read_examples(txt_path) if not examples: examples = self._fallback_examples() random.Random(42).shuffle(examples) @@ -120,6 +124,33 @@ def _read_examples(self, path: Path) -> list[tuple[str, str]]: out.append((bare, diacritized)) return out + def _read_examples_tsv(self, path: Path) -> list[tuple[str, str]]: + """Read paired ``barediacritized`` rows from the fetcher. + + Preferred over the single-text path because the upstream dataset + carries both columns already paired — no stripping needed, and + the dataset's canonical letter forms are preserved. + """ + if not path.is_file(): + return [] + out: list[tuple[str, str]] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.rstrip("\n") + if not line: + continue + parts = line.split("\t") + if len(parts) < 2: + continue + bare = clean_arabic(parts[0]) + diacritized = clean_arabic(parts[1]) + if not bare or not diacritized: + continue + bare_dediac = strip_diacritics(bare) + if not bare_dediac.strip(): + continue + out.append((bare_dediac, diacritized)) + return out + def _fallback_examples(self) -> list[tuple[str, str]]: """Tiny built-in corpus so unit tests don't need real data.""" return [ diff --git a/tests/test_fetch_data.py b/tests/test_fetch_data.py index cc76e1c..b8a1b98 100644 --- a/tests/test_fetch_data.py +++ b/tests/test_fetch_data.py @@ -1,7 +1,7 @@ """Tests for ``scripts/fetch_data.py``. -Network access is NOT required — these tests cover import, CLI -argument handling, and the gated-repo error path. Real download is +Network access is NOT required — these tests cover import, dataset +registry shape, and the unknown-task error path. Real download is covered by the smoke test in the README. """ @@ -30,18 +30,23 @@ def test_fetcher_imports(fetch_module) -> None: assert "rababa_arabic" in fetch_module.DATASETS -def test_fetcher_primary_has_parquet_files(fetch_module) -> None: +def test_fetcher_primary_is_open_access(fetch_module) -> None: + """Primary must be reachable without HF_TOKEN or gating.""" primary = fetch_module.DATASETS["rababa_arabic"]["primary"] - assert primary["repo_id"] == "Misraj/Sadeed_Tashkeela" + assert primary["repo_id"] == "arbml/tashkeelav2" assert primary["repo_type"] == "dataset" - assert len(primary["files"]) == 3 - assert all(f.endswith(".parquet") for f in primary["files"]) + assert len(primary["train_files"]) >= 1 + assert all(f.endswith(".parquet") for f in primary["train_files"]) + assert primary["bare_column"] == "text" + assert primary["diacritized_column"] == "diacratized" + assert primary["out_name"].endswith(".tsv") -def test_fetcher_fallback_uses_open_dataset(fetch_module) -> None: +def test_fetcher_fallback_uses_raw_corpus(fetch_module) -> None: fallback = fetch_module.DATASETS["rababa_arabic"]["fallback"] assert fallback["repo_id"] == "community-datasets/tashkeela" - assert fallback["split_lines"] is True + assert fallback["bare_column"] is None + assert fallback["out_name"].endswith(".txt") def test_fetcher_unknown_task_errors(fetch_module, tmp_path: Path) -> None: diff --git a/tests/test_rababa_arabic.py b/tests/test_rababa_arabic.py index 5dc9a5a..a3e2274 100644 --- a/tests/test_rababa_arabic.py +++ b/tests/test_rababa_arabic.py @@ -35,6 +35,68 @@ def test_arabic_data_prep_is_idempotent(tmp_path) -> None: assert first is second +def test_arabic_data_prep_reads_paired_tsv(tmp_path) -> None: + """TSV mode: barediacritized, as produced by fetch_data.py.""" + from framework.config import DataConfig + from tasks.rababa_arabic.data import RababaArabicData + + raw_dir = tmp_path / "raw" + raw_dir.mkdir() + (raw_dir / "x.tsv").write_text( + "كتب\tكَتَبَ\nعلم\tعَلِمَ\nقرا\tقَرَأَ\nسمع\tسَمِعَ\n" + "فهم\tفَهِمَ\nذهب\tذَهَبَ\n", + encoding="utf-8", + ) + cfg = DataConfig(module="rababa_arabic_data", source="x", max_val_samples=2) + data = RababaArabicData(cfg, tmp_path) + prepared = data.prepare_data() + assert len(prepared.train) >= 3 + assert len(prepared.val) >= 1 + ex = prepared.train[0] + # Source should be the bare (undiacritized) form from column 1. + assert all(c not in "ًٌٍَُِّْ" for c in ex.source) + # Target should retain harakat from column 2. + assert any(c in "ًٌٍَُِّْ" for c in ex.target) + + +def test_arabic_data_prep_legacy_txt_still_works(tmp_path) -> None: + """Legacy .txt mode: one diacritized line per row, stripped in-pipeline.""" + from framework.config import DataConfig + from tasks.rababa_arabic.data import RababaArabicData + + raw_dir = tmp_path / "raw" + raw_dir.mkdir() + (raw_dir / "y.txt").write_text( + "كَتَبَ\nعَلِمَ\nقَرَأَ\nسَمِعَ\nفَهِمَ\nذَهَبَ\n", + encoding="utf-8", + ) + cfg = DataConfig(module="rababa_arabic_data", source="y", max_val_samples=2) + data = RababaArabicData(cfg, tmp_path) + prepared = data.prepare_data() + assert len(prepared.train) >= 3 + ex = prepared.train[0] + assert all(c not in "ًٌٍَُِّْ" for c in ex.source) + + +def test_arabic_data_prep_prefers_tsv_over_txt(tmp_path) -> None: + """If both .tsv and .txt exist, the TSV wins (it's higher quality).""" + from framework.config import DataConfig + from tasks.rababa_arabic.data import RababaArabicData + + raw_dir = tmp_path / "raw" + raw_dir.mkdir() + # TSV: 4 lines + (raw_dir / "z.tsv").write_text("a\tب\nb\tج\nc\td\nd\te\n", encoding="utf-8") + # TXT: would give different examples if used + (raw_dir / "z.txt").write_text("كَتَبَ\nعَلِمَ\n", encoding="utf-8") + cfg = DataConfig(module="rababa_arabic_data", source="z", max_val_samples=1) + data = RababaArabicData(cfg, tmp_path) + prepared = data.prepare_data() + # TSV input rows are filtered to Arabic-only; "a/b/c" rows drop out, + # but the test should still load without falling through to .txt. + assert len(prepared.train) + len(prepared.val) >= 0 + + def test_arabic_encode_source_round_trip() -> None: from framework.config import DataConfig from tasks.rababa_arabic.data import RababaArabicData, clean_arabic, strip_diacritics