scripts/docker_build_assets.py

#!/usr/bin/env python3

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 _local_embedding_model() -> str:
    return (
        os.environ.get("LOCAL_EMBEDDING_MODEL", "").strip()
        or os.environ.get("TASK_B_LOCAL_EMBEDDING_MODEL", "").strip()
        or os.environ.get("TASK_A_EMBEDDING_MODEL", "").strip()
        or "all-MiniLM-L6-v2"
    )


def _skip_local_llm_hub() -> bool:
    if os.environ.get("SKIP_LOCAL_LLM_HUB_DOWNLOAD", "").strip().lower() in (
        "1",
        "true",
        "yes",
    ):
        return True
    if os.environ.get("GENERATION_BACKEND", "").strip().lower() in ("gemini", "google"):
        return True
    return False


def _unique_llm_hub_ids() -> list[str]:
    if _skip_local_llm_hub():
        return []
    seen: list[str] = []
    for key in ("LOCAL_LLM_MODEL", "TASK_B_LOCAL_LLM_MODEL", "TASK_A_LOCAL_LLM_MODEL"):
        v = os.environ.get(key, "").strip()
        if v and v not in seen:
            seen.append(v)
    if not seen:
        seen.append("Qwen/Qwen2.5-1.5B-Instruct")
    return seen


_build_embedder: Any = None


def _get_build_embedder() -> Any:
    global _build_embedder
    if _build_embedder is None:
        from sentence_transformers import SentenceTransformer  # type: ignore[import-untyped]

        emb_name = _local_embedding_model()
        print(f"docker_build_assets: loading embedder once for stub JSONL -> {emb_name}")
        _build_embedder = SentenceTransformer(emb_name)
    return _build_embedder


def prefetch_hub_files_only() -> None:
    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(_local_embedding_model())
    print(f"docker_build_assets: snapshot_download (disk cache) -> {emb_repo}")
    snapshot_download(repo_id=emb_repo, local_files_only=False, **kw)

    llm_ids = _unique_llm_hub_ids()
    if not llm_ids:
        print(
            "docker_build_assets: skipping causal LM snapshot "
            "(GENERATION_BACKEND=gemini or SKIP_LOCAL_LLM_HUB_DOWNLOAD=1)."
        )
    for llm_id in llm_ids:
        print(f"docker_build_assets: snapshot_download (disk cache) -> {llm_id}")
        snapshot_download(repo_id=llm_id, local_files_only=False, **kw)

    print("docker_build_assets: Hub snapshots cached (embedder on disk; LLM via API or runtime load).")


def warm_runtime_models() -> None:
    raw = os.environ.get("DOCKER_BUILD_SKIP_LLM_WARM", "1").strip().lower()
    skip = raw not in ("0", "false", "no")
    if skip:
        print(
            "docker_build_assets: skipping in-RAM LLM warm (DOCKER_BUILD_SKIP_LLM_WARM default 1). "
            "Weights are on disk from snapshot_download + stub encodes; uvicorn prewarm loads them at runtime."
        )
        return

    print("docker_build_assets: full model warm (CPU) — DOCKER_BUILD_SKIP_LLM_WARM=0; needs several GB RAM.")
    import gc

    st = _get_build_embedder()
    st.encode(["docker-build-warmup"], batch_size=1, show_progress_bar=False, convert_to_numpy=True)
    gc.collect()

    import torch  # type: ignore[import-untyped]
    from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore[import-untyped]

    for mid in _unique_llm_hub_ids():
        print(f"docker_build_assets: causal LM warm — {mid}")
        tok = AutoTokenizer.from_pretrained(mid, trust_remote_code=True)
        mdl = AutoModelForCausalLM.from_pretrained(
            mid,
            torch_dtype=torch.float32,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        mdl.eval()
        with torch.no_grad():
            batch = tok("warmup", return_tensors="pt")
            mdl(**batch)
        del mdl, tok
        gc.collect()

    print("docker_build_assets: model warm complete.")


def yelp_review_path(rt: Path) -> Path:
    env_p = os.environ.get("YELP_REVIEW_JSON", "").strip()
    if env_p:
        return Path(env_p)
    return rt / "yelp_dataset" / "extracted" / "yelp_academic_dataset_review.json"


def build_task_a_from_yelp(rt: Path, business_json: Path, review_json: Path) -> None:
    py = sys.executable
    max_rev = os.environ.get("DOCKER_TASK_A_MAX_REVIEW_ROWS", "10000")
    subprocess.check_call(
        [
            py,
            str(rt / "scripts" / "build_task_a_review_rag.py"),
            "--review-json",
            str(review_json),
            "--business-json",
            str(business_json),
            "--output",
            str(rt / "data" / "task_a_reviews_embedded.jsonl"),
            "--max-rows",
            max_rev,
            "--batch-size",
            "32",
        ]
    )


def build_stub_task_a_embedded(rt: Path, model: Any) -> None:
    rows_raw: list[dict[str, Any]] = []
    for i in range(36):
        uid = f"stub_user_{i % 9}"
        excerpt = (
            f"Stub visit {i}: food was {'solid' if i % 2 == 0 else 'hit or miss'}, "
            f"service {'friendly' if i % 3 else 'slow'}."
        )
        bctx = (
            f"name: Stub Diner {i % 6}\ncategories: Restaurants, "
            f"{'Italian' if i % 3 else 'Mexican'}\nlocation: Lagos, LA\n"
            f"business_avg_stars: {3.4 + (i % 8) / 10:.1f}\n"
            f"business_review_count: {15 + i * 3}"
        )
        emb_src = f"{bctx}\nreview: {excerpt}"
        rows_raw.append(
            {
                "user_id": uid,
                "business_id": f"stub_b_{i}",
                "stars": 1 + (i % 5),
                "review_excerpt": excerpt,
                "business_context": bctx,
                "embedding_source": emb_src,
            }
        )
    texts = [r["embedding_source"] for r in rows_raw]
    mat = model.encode(texts, batch_size=8, convert_to_numpy=True, normalize_embeddings=False)
    out_path = rt / "data" / "task_a_reviews_embedded.jsonl"
    with out_path.open("w", encoding="utf-8") as fout:
        for rec, vec in zip(rows_raw, mat, strict=True):
            row_out = {k: v for k, v in rec.items() if k != "embedding_source"}
            row_out["embedding"] = vec.astype(float).tolist()
            fout.write(json.dumps(row_out, ensure_ascii=False) + "\n")
    print(f"docker_build_assets: wrote stub Task A RAG -> {out_path}")


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.py"),
            "--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, model: Any) -> None:
    rows = stub_catalog_rows()
    texts = [r["text_for_embedding"] for r in rows]
    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 build_stub_data(rt: Path) -> None:
    model = _get_build_embedder()
    build_stub_embedded(rt, model)
    build_stub_task_a_embedded(rt, model)


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

    hf_login()
    prefetch_hub_files_only()

    yelp = yelp_business_path(rt)
    yelp_rev = yelp_review_path(rt)
    if yelp.is_file():
        print(f"docker_build_assets: building catalog from {yelp}")
        build_from_yelp(rt, yelp)
        if yelp_rev.is_file():
            print(f"docker_build_assets: Task A RAG from {yelp_rev}")
            build_task_a_from_yelp(rt, yelp, yelp_rev)
        else:
            print(
                "docker_build_assets: Yelp review JSON not found — stub Task A RAG "
                "(set YELP_REVIEW_JSON or add yelp_academic_dataset_review.json)."
            )
            build_stub_task_a_embedded(rt, _get_build_embedder())
    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_data(rt)

    warm_runtime_models()


if __name__ == "__main__":
    main()