"""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 /sources// -- 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