code.py

"""Week-2 20% mixed-bit base plus grouped-JSD residual entry points."""

from __future__ import annotations

import importlib.util
import json
import os
import shutil
import sys
import sysconfig
import tempfile
from pathlib import Path

# The course requires this filename, which otherwise shadows Python's stdlib
# ``code`` module while torch imports pdb. Export the stdlib API first.
if __name__ == "code":
    _stdlib_path = Path(sysconfig.get_path("stdlib")) / "code.py"
    _spec = importlib.util.spec_from_file_location("_cs6013_stdlib_code", _stdlib_path)
    if _spec is None or _spec.loader is None:
        raise ImportError(f"could not load stdlib code module: {_stdlib_path}")
    _stdlib = importlib.util.module_from_spec(_spec)
    _spec.loader.exec_module(_stdlib)
    for _name in ("InteractiveInterpreter", "InteractiveConsole", "interact", "compile_command"):
        globals()[_name] = getattr(_stdlib, _name)

import torch

LOCAL_SRC = Path(__file__).resolve().parent / "src"
if LOCAL_SRC.is_dir() and str(LOCAL_SRC) not in sys.path:
    sys.path.insert(0, str(LOCAL_SRC))

from eaimath.artifact import (  # noqa: E402
    dequantize_gptq_model,
    pack_model_state,
    read_trace_corpus,
    restore_artifact,
    save_artifact,
)
from eaimath.buckets import BY_NAME  # noqa: E402
from eaimath.embedding_predictor import fit_token_predictor  # noqa: E402
from eaimath.model import load_model, load_tokenizer, save_checkpoint  # noqa: E402
from eaimath.pack import state_dict_bytes  # noqa: E402
from eaimath.peft_compat import prepare_dense_lora_dispatch  # noqa: E402
from eaimath.vocab import build_keep_set, token_frequencies  # noqa: E402

SUBMISSION_HF_REPO = "safffrron/25M2111-Week02-Track2-20-Submission01"
BASE_FILENAME = "week02_20_base.ptz"
KEEP_FILENAME = "week02_20_keep_ids.json"
ADAPTER_DIRNAME = "adapter"
EXPECTED_BITS = {"mlp": 3, "linear_attn": 4, "full_attn": 8, "embed": 8}


def _reproduction_inputs() -> tuple[Path, Path, Path, dict]:
    source = Path(os.environ.get("EAIMATH_B20_GPTQ_SOURCE", "work/gptq_m3l4a8e8"))
    corpus = Path(os.environ.get("EAIMATH_VOCAB_CORPUS", "data/traces.jsonl"))
    adapter = Path(os.environ.get("EAIMATH_WEEK2_ADAPTER_SOURCE", "work/b_postq_jsd_r8"))
    required = [source / "experiment_config.json", corpus, adapter / "adapter_model.safetensors", adapter / "adapter_config.json"]
    missing = [str(path) for path in required if not path.is_file()]
    if missing:
        raise FileNotFoundError(
            "Reproduction inputs are missing. Follow README.md's rebuild section or set "
            "EAIMATH_B20_GPTQ_SOURCE, EAIMATH_VOCAB_CORPUS, and "
            f"EAIMATH_WEEK2_ADAPTER_SOURCE. Missing: {missing}"
        )
    config = json.loads((source / "experiment_config.json").read_text())
    if config.get("bits") != EXPECTED_BITS:
        raise ValueError(f"expected m3l4a8e8, found bits={config.get('bits')}")
    if set(config.get("gptq_components", [])) != {"mlp", "linear_attn"}:
        raise ValueError("GPTQ source must calibrate both MLP and linear attention")
    if config.get("activation_order") != "gar":
        raise ValueError("GPTQ source must use GAR activation order")
    adapter_config = json.loads((adapter / "adapter_config.json").read_text())
    if int(adapter_config.get("r", -1)) != 8:
        raise ValueError("the submitted residual must have rank 8")
    return source, corpus, adapter, config


def _locate_payload(checkpoint_path: str) -> tuple[Path, Path, Path]:
    supplied = Path(checkpoint_path).expanduser().resolve()
    root = supplied if supplied.is_dir() else supplied.parent
    base = root / BASE_FILENAME if supplied.is_dir() else supplied
    keep = root / KEEP_FILENAME
    adapter = root / ADAPTER_DIRNAME
    if base.is_file() and keep.is_file() and (adapter / "adapter_model.safetensors").is_file():
        return base, keep, adapter

    from huggingface_hub import snapshot_download

    snapshot = Path(
        snapshot_download(
            SUBMISSION_HF_REPO,
            allow_patterns=[BASE_FILENAME, f"{BASE_FILENAME}.json", KEEP_FILENAME, "adapter/*"],
        )
    )
    base, keep, adapter = snapshot / BASE_FILENAME, snapshot / KEEP_FILENAME, snapshot / ADAPTER_DIRNAME
    missing = [str(path) for path in (base, keep, adapter / "adapter_model.safetensors", adapter / "adapter_config.json") if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"incomplete Week-2 20% payload: {missing}")
    return base, keep, adapter


def convert_from_hf_checkpoint(
    model_name: str,
    output_path: str,
    sparsity: float | None = None,
) -> None:
    """Reproduce the packed base and copy the trained JSD residual beside it.

    ``sparsity`` is accepted only because the starter interface may pass it;
    this method is dense mixed-bit quantization and does not use that value.
    """
    _ = sparsity
    source, corpus, adapter_source, source_config = _reproduction_inputs()
    requested = Path(output_path)
    if requested.suffix:
        root, artifact = requested.parent, requested
    else:
        root, artifact = requested, requested / BASE_FILENAME
    root.mkdir(parents=True, exist_ok=True)
    if any((root / name).exists() for name in (KEEP_FILENAME, ADAPTER_DIRNAME)):
        raise FileExistsError(f"refusing to overwrite an existing representation in {root}")

    problems, completions, corpus_stats = read_trace_corpus([corpus])
    tokenizer = load_tokenizer(str(source))
    counts = token_frequencies(problems + completions, tokenizer)
    priority = token_frequencies(problems, tokenizer)
    keep_set = build_keep_set(counts, tokenizer, 30_000, priority_counts=priority)
    loaded, model, replaced, exact_weights, capture_error = dequantize_gptq_model(source)
    if capture_error != 0.0:
        raise RuntimeError(f"GPTQ code capture error: {capture_error}")
    payload, _ = pack_model_state(
        model.state_dict(), keep_set["keep_ids"], group_size=128,
        exact_entries=exact_weights, recipe_bits=EXPECTED_BITS, recipe_name="m3l4a8e8",
    )
    payload["source"] = {
        "base_model": model_name,
        "experiment_config": source_config,
        "vocab_strategy": "problem-first",
        "corpus_stats": corpus_stats,
        "gptq_linears": replaced,
    }
    row_names = [name for name, entry in payload["state_dict"].items() if entry["kind"] in {"quant_rows", "raw_rows"}]
    if len(row_names) != 1:
        raise RuntimeError(f"expected one canonical tied embedding, found {row_names}")
    predictor, predictor_report = fit_token_predictor(
        model.state_dict()[row_names[0]].detach().cpu(), tokenizer,
        sample_size=65_536, device=os.environ.get("EAIMATH_PREDICTOR_DEVICE", "cpu"),
    )
    payload["embedding_predictor"] = predictor
    payload["report"]["embedding_predictor"] = predictor_report
    payload["report"]["tensor_bytes"] += state_dict_bytes({"embedding_predictor": predictor["basis"]})
    (root / KEEP_FILENAME).write_text(json.dumps(keep_set["keep_ids"]))
    final = save_artifact(payload, artifact, lossless_codec="zlib", zlib_level=6)
    shutil.copytree(adapter_source, root / ADAPTER_DIRNAME, ignore=shutil.ignore_patterns("__pycache__", "*.pyc"))
    total = int(final["disk_bytes"]) + (root / KEEP_FILENAME).stat().st_size + sum(
        path.stat().st_size for path in (root / ADAPTER_DIRNAME).rglob("*") if path.is_file()
    )
    ceiling = int(BY_NAME["B"].high_gb * 1e9)
    if total >= ceiling:
        raise RuntimeError(f"full representation exceeds the 20% ceiling: {total:,} >= {ceiling:,}")
    del model, loaded
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def convert_to_hf_checkpoint(model_name: str, checkpoint_path: str, output_path: str) -> None:
    """Restore the packed base, safely merge the residual, and save BF16 HF."""
    from peft import PeftModel

    base, keep, adapter = _locate_payload(checkpoint_path)
    output = Path(output_path)
    if output.exists() and (not output.is_dir() or any(output.iterdir())):
        raise FileExistsError(f"refusing to overwrite non-empty output: {output}")
    output.parent.mkdir(parents=True, exist_ok=True)
    device = os.environ.get("EAIMATH_RESTORE_DEVICE", "cpu")
    with tempfile.TemporaryDirectory(prefix="week02-b20-", dir=output.parent) as temporary:
        restored_base = Path(temporary) / "base"
        base_report = restore_artifact(model_name, base, restored_base, embedding_fill="token")
        model = load_model(str(restored_base), dtype="bfloat16", device=device, multimodal=True)
        prepare_dense_lora_dispatch(model)
        merged = PeftModel.from_pretrained(model, adapter).merge_and_unload(safe_merge=True)
        merged.config.use_cache = True
        copied = save_checkpoint(merged, output, source_model=model_name)
    shutil.copy2(keep, output / "keep_ids.json")
    report = {
        "format": "eaimath-week02-b20-restored-v1",
        "base_artifact": str(base),
        "adapter": str(adapter),
        "safe_merge": True,
        "base_restore": base_report,
        "auxiliary_files": copied,
    }
    (output / "submission_report.json").write_text(json.dumps(report, indent=2) + "\n")