add cache

fleurs_cache.py

@@ -0,0 +1,323 @@
+from __future__ import annotations
+
+import argparse
+import random
+import unicodedata
+from functools import lru_cache
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+from huggingface_hub import HfApi, hf_hub_download
+
+from language import ALL_LANGS, LANG_ISO2_TO_ISO3, canonical_lang
+
+
+FLEURS_DATASET = "google/fleurs"
+FLEURS_CACHE_DIR = Path(__file__).with_name("data") / "fleurs"
+FLEURS_PARQUET_PATH = FLEURS_CACHE_DIR / "fleurs_text_only.parquet"
+FLEURS_DOWNLOAD_DIR = FLEURS_CACHE_DIR / "downloads"
+FLEURS_TSV_COLUMNS = [
+    "id",
+    "file_name",
+    "source_sentence",
+    "transcription",
+    "tokens",
+    "num_samples",
+    "gender",
+]
+FLEURS_SPLIT_ORDER = {"train": 0, "validation": 1, "test": 2}
+FLEURS_LEAN_COLUMNS = ["id", "text", "source_lang", "model_lang", "split"]
+
+
+def _normalize_model_lang(source_lang: str) -> str:
+    """Map a FLEURS locale like `am_et` to the model language code."""
+    base_lang = source_lang.split("_", 1)[0].strip().lower()
+    return canonical_lang(base_lang)
+
+
+def _discover_tsv_files() -> list[str]:
+    """Return all FLEURS TSV metadata files, preferring the local cache."""
+    local_root = FLEURS_DOWNLOAD_DIR / "data"
+    local_files = sorted(local_root.rglob("*.tsv"))
+    if local_files:
+        return [str(path.relative_to(FLEURS_DOWNLOAD_DIR)) for path in local_files]
+
+    api = HfApi()
+    try:
+        files = api.list_repo_files(repo_id=FLEURS_DATASET, repo_type="dataset")
+    except TypeError:
+        files = api.list_repo_files(FLEURS_DATASET, repo_type="dataset")
+
+    tsv_files = [
+        file_path
+        for file_path in files
+        if file_path.startswith("data/") and file_path.endswith(".tsv")
+    ]
+    if not tsv_files:
+        raise RuntimeError("Could not find any FLEURS TSV metadata files.")
+    return sorted(tsv_files)
+
+
+def _normalize_split_name(file_name: str) -> str:
+    stem = Path(file_name).stem.lower()
+    if stem == "dev":
+        return "validation"
+    return stem
+
+
+def _normalize_text_key(text: str) -> str:
+    """Normalize text for deduping while keeping the original text intact."""
+    normalized = unicodedata.normalize("NFKC", text)
+    normalized = " ".join(normalized.split())
+    return normalized.casefold().strip()
+
+
+def _download_tsv(file_path: str) -> Path:
+    local_candidate = FLEURS_DOWNLOAD_DIR / file_path
+    if local_candidate.exists():
+        return local_candidate
+
+    FLEURS_DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)
+    try:
+        local_path = hf_hub_download(
+            repo_id=FLEURS_DATASET,
+            repo_type="dataset",
+            filename=file_path,
+            local_dir=str(FLEURS_DOWNLOAD_DIR),
+        )
+    except TypeError:
+        local_path = hf_hub_download(
+            repo_id=FLEURS_DATASET,
+            repo_type="dataset",
+            filename=file_path,
+            cache_dir=str(FLEURS_DOWNLOAD_DIR),
+        )
+    return Path(local_path)
+
+
+def _frame_from_tsv(tsv_path: Path, source_lang: str) -> pd.DataFrame:
+    records: list[dict[str, Any]] = []
+    header_seen = False
+    header_markers = {name.lower() for name in FLEURS_TSV_COLUMNS}
+
+    with tsv_path.open("r", encoding="utf-8", newline="") as handle:
+        for line in handle:
+            line = line.rstrip("\n")
+            if not line.strip():
+                continue
+
+            parts = line.split("\t", 6)
+            if not header_seen:
+                header_candidate = [part.strip().lower() for part in parts]
+                if header_markers.issubset(set(header_candidate)):
+                    header_seen = True
+                    continue
+                header_seen = True
+
+            if len(parts) < len(FLEURS_TSV_COLUMNS):
+                parts.extend([""] * (len(FLEURS_TSV_COLUMNS) - len(parts)))
+            elif len(parts) > len(FLEURS_TSV_COLUMNS):
+                parts = parts[: len(FLEURS_TSV_COLUMNS) - 1] + ["\t".join(parts[len(FLEURS_TSV_COLUMNS) - 1 :])]
+
+            record = dict(zip(FLEURS_TSV_COLUMNS, parts, strict=True))
+            records.append(record)
+
+    if not records:
+        return pd.DataFrame()
+
+    frame = pd.DataFrame.from_records(records)
+    frame = frame.fillna("")
+    frame["source_sentence"] = frame["source_sentence"].astype(str).str.strip()
+    frame["transcription"] = frame["transcription"].astype(str).str.strip()
+    frame["tokens"] = frame["tokens"].astype(str).str.strip()
+
+    frame["text"] = frame["transcription"].where(frame["transcription"].ne(""), frame["source_sentence"])
+    frame["raw_text"] = frame["source_sentence"].where(frame["source_sentence"].ne(""), frame["transcription"])
+    frame["source"] = "fleurs"
+    frame["source_lang"] = source_lang
+    frame["model_lang"] = _normalize_model_lang(source_lang)
+    frame["split"] = _normalize_split_name(tsv_path.name)
+    frame["lang_iso3"] = frame["model_lang"].map(lambda lang: LANG_ISO2_TO_ISO3.get(lang, ""))
+    frame["language_name"] = source_lang
+    frame["text"] = frame["text"].astype(str).str.strip().replace("", pd.NA)
+    frame["raw_text"] = frame["raw_text"].astype(str).str.strip()
+    frame = frame[frame["text"].notna()].reset_index(drop=True)
+    return frame
+
+
+def _post_process_fleurs_frame(frame: pd.DataFrame) -> pd.DataFrame:
+    """Drop redundant rows and keep only the lean demo columns."""
+    if frame.empty:
+        return frame
+
+    frame = frame.copy()
+    frame["split_rank"] = frame["split"].map(lambda split: FLEURS_SPLIT_ORDER.get(str(split), 99))
+    frame["text_key"] = frame["text"].astype(str).map(_normalize_text_key)
+    frame["id_sort"] = pd.to_numeric(frame["id"], errors="coerce").fillna(10**18)
+
+    frame = frame[frame["text_key"].ne("")].sort_values(
+        by=["source_lang", "text_key", "split_rank", "id_sort"],
+        kind="stable",
+    )
+    frame = frame.drop_duplicates(subset=["source_lang", "text_key"], keep="first")
+
+    lean = frame.loc[:, [col for col in FLEURS_LEAN_COLUMNS if col in frame.columns]].copy()
+    lean["text"] = frame["text"].astype(str).values
+    lean["source_lang"] = frame["source_lang"].astype(str).values
+    lean["model_lang"] = frame["model_lang"].astype(str).values
+    lean["split"] = frame["split"].astype(str).values
+    lean["id"] = pd.to_numeric(frame["id"], errors="coerce").fillna(-1).astype(int).values
+    lean = lean[lean["text"].astype(str).str.strip().ne("")].reset_index(drop=True)
+    return lean
+
+
+def build_fleurs_text_parquet(
+    parquet_path: str | Path = FLEURS_PARQUET_PATH,
+) -> Path:
+    """Download FLEURS TSV metadata and persist a text-only parquet cache."""
+    parquet_path = Path(parquet_path)
+    parquet_path.parent.mkdir(parents=True, exist_ok=True)
+
+    frames: list[pd.DataFrame] = []
+    for repo_path in _discover_tsv_files():
+        source_lang = Path(repo_path).parent.name
+        tsv_path = _download_tsv(repo_path)
+        frame = _frame_from_tsv(tsv_path, source_lang)
+        if not frame.empty:
+            frames.append(frame)
+
+    if not frames:
+        raise RuntimeError("No rows were loaded from the FLEURS TSV metadata files.")
+
+    combined = pd.concat(frames, ignore_index=True)
+    before_rows = len(combined)
+    combined = _post_process_fleurs_frame(combined)
+    combined.to_parquet(parquet_path, index=False)
+    print(
+        f"Built lean FLEURS parquet with {len(combined):,} rows "
+        f"from {before_rows:,} raw rows and {len(combined.columns)} columns."
+    )
+    return parquet_path
+
+
+@lru_cache(maxsize=1)
+def load_fleurs_table(parquet_path: str | Path = FLEURS_PARQUET_PATH) -> pd.DataFrame:
+    """Load the cached FLEURS text-only parquet into memory."""
+    parquet_path = Path(parquet_path)
+    if not parquet_path.exists():
+        raise FileNotFoundError(
+            f"Missing FLEURS cache at {parquet_path}. "
+            "Run `./.venv/bin/python fleurs_cache.py` once while online to build it."
+        )
+
+    frame = pd.read_parquet(parquet_path)
+    if "text" not in frame.columns:
+        raise RuntimeError("FLEURS parquet cache is missing the text column.")
+    return frame
+
+
+def _pick_random_rows(frame: pd.DataFrame, *, count: int) -> pd.DataFrame:
+    if frame.empty:
+        raise RuntimeError("FLEURS cache has no rows.")
+    sample_size = min(count, len(frame))
+    return frame.sample(n=sample_size)
+
+
+def _row_to_sentence(row: pd.Series) -> dict[str, Any]:
+    source_lang = str(row.get("source_lang", "")).strip()
+    model_lang = str(row.get("model_lang", "")).strip()
+    lang_iso2 = model_lang or _normalize_model_lang(source_lang)
+    language = str(row.get("language_name", source_lang)).strip()
+    text = str(row.get("text", "")).strip()
+    return {
+        "text": text,
+        "raw_text": text,
+        "source": "fleurs",
+        "source_lang": source_lang,
+        "model_lang": model_lang or lang_iso2,
+        "lang_iso2": lang_iso2,
+        "lang_iso3": LANG_ISO2_TO_ISO3.get(lang_iso2, ""),
+        "language": language,
+        "split": str(row.get("split", "")).strip(),
+        "fleurs_id": int(row.get("id", -1)) if str(row.get("id", "-1")).strip().lstrip("-").isdigit() else -1,
+    }
+
+
+def fetch_random_fleurs_sentence(
+    *,
+    attempts: int = 8,
+    parquet_path: str | Path = FLEURS_PARQUET_PATH,
+) -> dict[str, Any]:
+    """Fetch one random text sample from the cached FLEURS parquet."""
+    frame = load_fleurs_table(parquet_path)
+    candidate_frame = frame[frame["text"].astype(str).str.strip().ne("")]
+
+    supported = candidate_frame[candidate_frame["model_lang"].isin(ALL_LANGS)]
+    if not supported.empty:
+        candidate_frame = supported
+
+    for _ in range(max(1, attempts)):
+        row = _pick_random_rows(candidate_frame, count=1).iloc[0]
+        sentence = _row_to_sentence(row)
+        if sentence["text"]:
+            return sentence
+
+    raise RuntimeError("Unable to sample a random FLEURS sentence.")
+
+
+def fetch_random_fleurs_sentence_mix(
+    *,
+    min_sentences: int = 2,
+    max_sentences: int = 3,
+    parquet_path: str | Path = FLEURS_PARQUET_PATH,
+) -> dict[str, Any]:
+    """Fetch 2-3 random FLEURS sentences from distinct languages and concatenate them."""
+    frame = load_fleurs_table(parquet_path)
+    candidate_frame = frame[frame["text"].astype(str).str.strip().ne("")]
+    supported = candidate_frame[candidate_frame["model_lang"].isin(ALL_LANGS)]
+    if not supported.empty:
+        candidate_frame = supported
+
+    min_sentences = max(1, min_sentences)
+    max_sentences = max(min_sentences, max_sentences)
+    count = random.randint(min_sentences, max_sentences)
+
+    distinct_langs = [lang for lang in candidate_frame["model_lang"].dropna().unique().tolist() if lang]
+    if not distinct_langs:
+        raise RuntimeError("No usable FLEURS languages were found in the cache.")
+
+    random.shuffle(distinct_langs)
+    chosen_langs = distinct_langs[: min(count, len(distinct_langs))]
+
+    rows = []
+    for lang in chosen_langs:
+        lang_rows = candidate_frame[candidate_frame["model_lang"] == lang]
+        rows.append(_pick_random_rows(lang_rows, count=1).iloc[0])
+
+    sentences = [_row_to_sentence(row) for row in rows]
+    combined_text = "\n\n".join(sentence["text"] for sentence in sentences if sentence["text"])
+    return {
+        "text": combined_text,
+        "sentences": sentences,
+        "lang_count": len(sentences),
+        "langs": [sentence["lang_iso2"] for sentence in sentences],
+        "lang_iso3s": [sentence["lang_iso3"] for sentence in sentences],
+        "source": "fleurs-mix",
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Build the cached text-only FLEURS parquet.")
+    parser.add_argument(
+        "--output",
+        default=str(FLEURS_PARQUET_PATH),
+        help="Output parquet path for the cached FLEURS text rows.",
+    )
+    args = parser.parse_args()
+    path = build_fleurs_text_parquet(args.output)
+    print(f"Wrote FLEURS text cache to {path}")
+
+
+if __name__ == "__main__":
+    main()