scripts/docker_build_assets.py

#!/usr/bin/env python3
"""Docker build: HF auth, cache model files on disk (low RAM), produce data/*.jsonl."""

from __future__ import annotations

import hashlib
import json
import os
import subprocess
import sys
from pathlib import Path
from typing import Any


def root() -> Path:
    return Path(__file__).resolve().parents[1]


def hf_token() -> str | None:
    t = (
        os.environ.get("HF_TOKEN", "").strip()
        or os.environ.get("HUGGING_FACE_HUB_TOKEN", "").strip()
    )
    return t or None


def hf_login() -> None:
    tok = hf_token()
    if not tok:
        print(
            "docker_build_assets: No HF_TOKEN / HUGGING_FACE_HUB_TOKEN — anonymous Hub access "
            "(rate limits). On HF Spaces, pass token into the *Docker build* (not only runtime)."
        )
        return
    try:
        from huggingface_hub import login  # type: ignore[import-untyped]
    except ImportError:
        print("docker_build_assets: huggingface_hub not installed; skipping login.")
        return
    login(token=tok, add_to_git_credential=False)
    print("docker_build_assets: Hugging Face Hub login OK.")


def embedding_model_repo(hub_name: str) -> str:
    if "/" not in hub_name.strip():
        return f"sentence-transformers/{hub_name.strip()}"
    return hub_name.strip()


def prefetch_hub_files_only() -> None:
    """Download weights into HF cache without loading full LLM into RAM (avoids build OOM)."""
    emb_name = os.environ.get("TASK_B_LOCAL_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
    llm_name = os.environ.get("TASK_B_LOCAL_LLM_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
    tok = hf_token()

    try:
        from huggingface_hub import snapshot_download  # type: ignore[import-untyped]
    except ImportError:
        print("docker_build_assets: huggingface_hub missing; skipping prefetch.")
        return

    kw: dict[str, Any] = {}
    if tok:
        kw["token"] = tok

    emb_repo = embedding_model_repo(emb_name)
    print(f"docker_build_assets: snapshot_download (disk cache) -> {emb_repo}")
    snapshot_download(repo_id=emb_repo, local_files_only=False, **kw)

    print(f"docker_build_assets: snapshot_download (disk cache) -> {llm_name}")
    snapshot_download(repo_id=llm_name, local_files_only=False, **kw)
    print("docker_build_assets: Hub snapshots cached (LLM not loaded into RAM).")


def yelp_business_path(rt: Path) -> Path:
    env_p = os.environ.get("YELP_BUSINESS_JSON", "").strip()
    if env_p:
        return Path(env_p)
    return rt / "yelp_dataset" / "extracted" / "yelp_academic_dataset_business.json"


def build_from_yelp(rt: Path, yelp: Path) -> None:
    max_rows = os.environ.get("DOCKER_CATALOG_MAX_ROWS", "15000")
    out_cat = rt / "data" / "business_catalog.jsonl"
    py = sys.executable
    subprocess.check_call(
        [
            py,
            str(rt / "scripts" / "build_business_catalog.py"),
            "--business-json",
            str(yelp),
            "--output",
            str(out_cat),
            "--max-rows",
            max_rows,
            "--only-open",
        ]
    )
    subprocess.check_call(
        [
            py,
            str(rt / "scripts" / "embed_catalog_azure_openai.py"),
            "--backend",
            "local",
            "--input",
            str(out_cat),
            "--output",
            str(rt / "data" / "business_catalog_embedded.jsonl"),
            "--batch-size",
            "32",
        ]
    )


def stub_catalog_rows(n: int = 48) -> list[dict[str, Any]]:
    templates = [
        ("Riverfront Ramen", "Restaurants, Japanese, Ramen", "Portland", "OR"),
        ("Oak Street Bakery", "Food, Bakeries, Coffee & Tea", "Austin", "TX"),
        ("Queen Vietnamese", "Restaurants, Vietnamese", "Philadelphia", "PA"),
        ("Campus Espresso", "Coffee & Tea, Cafes", "Seattle", "WA"),
        ("Park Yoga Studio", "Active Life, Yoga", "Denver", "CO"),
        ("Midtown Books", "Shopping, Books", "Chicago", "IL"),
        ("East Side Brewpub", "Nightlife, Breweries", "Milwaukee", "WI"),
        ("Family Thai Kitchen", "Restaurants, Thai", "Tempe", "AZ"),
        ("Uptown Nail Spa", "Beauty & Spas, Nail Salons", "Miami", "FL"),
        ("Lakeside Pizza", "Restaurants, Pizza", "Minneapolis", "MN"),
    ]
    rows = []
    for i in range(n):
        name, cats, city, state = templates[i % len(templates)]
        suffix = i // len(templates)
        disp = f"{name}" if suffix == 0 else f"{name} #{suffix}"
        h = hashlib.sha256(f"{i}-{disp}".encode()).hexdigest()[:22]
        bid = h
        text_for_embedding = (
            f"name: {disp}\n"
            f"categories: {cats}\n"
            f"location: {city}, {state}\n"
            f"address: {100 + i} Main St\n"
            f"business_avg_stars: {3.5 + (i % 15) / 10:.1f}\n"
            f"business_review_count: {20 + i * 7}\n"
            f"is_open: 1"
        )
        rows.append(
            {
                "business_id": bid,
                "name": disp,
                "categories": cats,
                "city": city,
                "state": state,
                "stars": float(3.5 + (i % 15) / 10),
                "review_count": int(20 + i * 7),
                "is_open": 1,
                "text_for_embedding": text_for_embedding,
            }
        )
    return rows


def build_stub_embedded(rt: Path) -> None:
    from sentence_transformers import SentenceTransformer  # type: ignore[import-untyped]

    emb_name = os.environ.get("TASK_B_LOCAL_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
    model = SentenceTransformer(emb_name)
    rows = stub_catalog_rows()
    texts = [r["text_for_embedding"] for r in rows]
    # Small batches keep peak RAM low on HF builders.
    mat = model.encode(texts, batch_size=8, convert_to_numpy=True, normalize_embeddings=False)
    out_path = rt / "data" / "business_catalog_embedded.jsonl"
    cat_path = rt / "data" / "business_catalog.jsonl"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with out_path.open("w", encoding="utf-8") as fe, cat_path.open("w", encoding="utf-8") as fc:
        for row, vec in zip(rows, mat, strict=True):
            fc.write(json.dumps(row, ensure_ascii=False) + "\n")
            emb_row = {**row, "embedding": vec.astype(float).tolist()}
            fe.write(json.dumps(emb_row, ensure_ascii=False) + "\n")
    print(f"docker_build_assets: wrote stub catalog -> {cat_path} and {out_path}")


def main() -> None:
    rt = root()
    (rt / "data").mkdir(parents=True, exist_ok=True)

    hf_login()
    prefetch_hub_files_only()

    yelp = yelp_business_path(rt)
    if yelp.is_file():
        print(f"docker_build_assets: building catalog from {yelp}")
        build_from_yelp(rt, yelp)
    else:
        print(
            "docker_build_assets: Yelp business JSON not found; "
            "writing stub JSONL (mount real data at runtime or bake yelp_dataset into build context)."
        )
        build_stub_embedded(rt)


if __name__ == "__main__":
    main()