hako/sources/loader.py

"""Frozen source model acquisition: streaming download -> decompose -> purge.

Proposition P-SRC (one-at-a-time streaming with purge). For source s:
    1. hf_hub_download fetches only model.onnx + model.onnx.data + tokenizer
       (token from env HF_TOKEN exclusively -- NEVER persisted to disk).
    2. Decomposition extracts the mandated geometry (int4_proto, int4_xi,
       keys, 4 adapters) into compact .npz artifacts.
    3. The raw blobs are deleted (MemoryManager.sweep_caches) since the
       decomposed twins exist; disk returns to ~artifact size.
Peak disk = max over s of |raw_s| + |artifacts|, not the sum -- the
one-at-a-time rule turns a 4.2 GiB requirement into ~1.9 GiB.
"""
from __future__ import annotations

import json
import logging
import os
import shutil
from pathlib import Path
from typing import Dict, List

log = logging.getLogger("hako.sources")

QWEN = "amd/Qwen2.5-0.5B-Instruct-quantized_int4-float16-cpu-onnx"
GRANITE = "amd/granite-4.0-1b-awq-grp32-oga-cpu"

MODEL_FILES = ["model.onnx", "model.onnx.data", "tokenizer.json",
               "tokenizer_config.json", "vocab.json", "merges.txt"]


def _token() -> str | None:
    """Read the HF token from the environment ONLY. Never log it."""
    return os.environ.get("HF_TOKEN") or None


def fetch_source(repo: str, cache_dir: Path, memmgr) -> Path:
    """Download one source's onnx + tokenizer files, then MATERIALIZE real
    (dereferenced) copies under <work>/sources/<name>/ -- onnx's external-
    data loader rejects HF-cache symlinks ('not regular file'). Returns the
    materialized snapshot dir (which the purge step deletes)."""
    import shutil as _shutil
    from huggingface_hub import hf_hub_download
    cache_dir = Path(cache_dir)
    cache_dir.mkdir(parents=True, exist_ok=True)
    mat = Path(memmgr.cfg.work_dir) / "sources" / repo.split("/")[-1]
    marker = mat / "model.onnx.data"
    if marker.exists():
        return mat                      # materialized copy survives purge
    need = 4.5 * 1024 ** 3 if "granite" in repo.lower() else 2.2 * 1024 ** 3
    memmgr.ensure_disk(need)
    local: str | None = None
    for fname in MODEL_FILES:
        try:
            p = hf_hub_download(repo_id=repo, filename=fname,
                                cache_dir=str(cache_dir), token=_token())
            local = local or str(Path(p).parent)   # snapshot dir
            log.info("fetched %s/%s", repo, fname)
        except Exception as exc:
            log.info("skip %s (%s)", fname, type(exc).__name__)
    if local is None:
        raise RuntimeError(f"could not fetch any file from {repo}")
    mat.mkdir(parents=True, exist_ok=True)
    for fname in MODEL_FILES:
        src = Path(local) / fname
        dst = mat / fname
        if src.exists() and not dst.exists():
            _shutil.copy(src, dst)   # dereferences symlinks
    return mat


def purge_source(repo: str, cache_dir: Path, memmgr) -> int:
    """Delete raw blobs + materialized copy AFTER decomposed artifacts exist."""
    import shutil as _shutil
    marker = Path("/home/z/my-project/download/hako/artifacts") / (
        repo.split("/")[-1] + ".decomposed.json")
    if not marker.exists():
        log.info("keep %s: no decomposed artifact yet", repo)
        return 0
    freed = 0
    mat = Path(memmgr.cfg.work_dir) / "sources" / repo.split("/")[-1]
    if mat.exists():
        freed += sum(f.stat().st_size for f in mat.rglob("*") if f.is_file())
        _shutil.rmtree(mat, ignore_errors=True)
    freed += memmgr.sweep_caches()
    return freed


def fetch_datasets_streaming(cfg, memmgr, max_docs: Dict[str, int]
                             ) -> Dict[str, List[str]]:
    """Stream EN + PT-BR + zh text. Returns {lang: [docs]}.

    Streaming (datasets) reads parquet row-groups over HTTP ranges, so the
    1.1 GiB shard of corpus-ptbr-v1 and the 467 MiB ultra-alpaca parquet are
    never fully materialized on disk (P-SRC applied to data).
    """
    out: Dict[str, List] = {"pt": [], "en": [], "zh": [], "pt_instr": []}
    from datasets import load_dataset

    # --- PT-BR 1: BrunoN-Dev/ultra-alpaca-ptbr (instruction pairs) --------
    try:
        ds = load_dataset("BrunoN-Dev/ultra-alpaca-ptbr", split="train",
                          streaming=True)
        n = 0
        for row in ds:
            # instruction format: prompt=[{role,content}], completion=str
            try:
                prompt = row["prompt"][0]["content"]
                resp = row["completion"]
                out["pt"].append((str(prompt) + "\n" + str(resp))[:2000])
                out["pt_instr"].append((str(prompt)[:512], str(resp)[:512]))
            except (KeyError, IndexError, TypeError):
                txt = " \n ".join(str(v) for v in row.values()
                                  if v and len(str(v)) > 1)
                out["pt"].append(txt[:2000])
            n += 1
            if n >= max_docs.get("pt_alpaca", 12000):
                break
        log.info("ultra-alpaca-ptbr streamed: %d docs (%d instruction pairs)",
                 n, len(out["pt_instr"]))
    except Exception as exc:
        log.warning("ultra-alpaca-ptbr failed: %s", exc)

    # --- PT-BR 2: BrunoN-Dev/corpus-ptbr-v1 (cached sample) ---------------
    # The shard has ONE 500k-row row group (~1.2 GiB decompressed) that
    # cannot be streamed incrementally; it is materialized by the isolated
    # corpus_cache subprocess (run before training, clean RAM) and read here
    # from the small JSONL cache.
    cache_file = Path(cfg.work_dir) / "cache" / "corpus_ptbr.jsonl"
    if cache_file.exists():
        with open(cache_file, "r", encoding="utf-8") as fh:
            for line in fh:
                try:
                    out["pt"].append(json.loads(line)["text"][:2000])
                except Exception:
                    continue
        log.info("corpus-ptbr-v1 from cache: %d docs", len(out["pt"]))

    # --- EN: wikitext-2-raw ----------------------------------------------
    try:
        ds = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1",
                          split="train", streaming=True)
        n = 0
        for row in ds:
            t = row.get("text", "").strip()
            if len(t) > 200:
                out["en"].append(t[:2000])
                n += 1
            if n >= max_docs.get("en", 8000):
                break
        log.info("wikitext-2 streamed: %d docs", n)
    except Exception as exc:
        log.warning("wikitext failed: %s", exc)

    # --- zh: wikipedia streaming slice ------------------------------------
    try:
        ds = load_dataset("wikimedia/wikipedia", "20231101.zh",
                          split="train", streaming=True)
        n = 0
        for row in ds:
            t = row.get("text", "").strip()
            if len(t) > 200:
                out["zh"].append(t[:1500])
                n += 1
            if n >= max_docs.get("zh", 2500):
                break
        log.info("wikipedia-zh streamed: %d docs", n)
    except Exception as exc:
        log.warning("wikipedia-zh failed: %s", exc)

    for lang in out:
        memmgr.cap_file_size  # no-op ref to keep linters calm
    return out