code.py

"""Week-1 Track-2 40% block-adaptive compression entry points.

The compressed checkpoint is self-contained.  Recreating it requires the
expanded Round-14 block64 source and its allocation report; restoration needs
only the compressed artifact plus the original base model identifier supplied
by the course interface.
"""

from __future__ import annotations

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

# This file name is required by the course interface, but ``code`` is also a
# Python standard-library module used by ``pdb`` during PyTorch import.  When a
# user runs a wrapper from this directory, Python can resolve this file for both
# names.  Publish the stdlib API before importing torch so that the recursive
# ``pdb -> code`` import remains valid.
if __name__ == "code":
    _stdlib_code_path = Path(sysconfig.get_path("stdlib")) / "code.py"
    _stdlib_code_spec = importlib.util.spec_from_file_location(
        "_cs6013_stdlib_code", _stdlib_code_path
    )
    if _stdlib_code_spec is None or _stdlib_code_spec.loader is None:
        raise ImportError(f"could not load Python stdlib code module: {_stdlib_code_path}")
    _stdlib_code = importlib.util.module_from_spec(_stdlib_code_spec)
    _stdlib_code_spec.loader.exec_module(_stdlib_code)
    for _stdlib_name in (
        "InteractiveInterpreter",
        "InteractiveConsole",
        "interact",
        "compile_command",
    ):
        globals()[_stdlib_name] = getattr(_stdlib_code, _stdlib_name)

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.adaptive_artifact import (
    pack_block_adaptive_state,
    restore_block_adaptive_artifact,
    save_block_adaptive_artifact,
)
from eaimath.model import load_model

SUBMISSION_HF_REPO = "safffrron/25M2111-Week01-Track2-40-Submission01"


def _allocation_path(source: str) -> Path:
    configured = os.environ.get("EAIMATH_BLOCK64_REPORT")
    if configured:
        path = Path(configured)
    elif Path(__file__).with_name("configs").joinpath("block_adaptive_report.json").is_file():
        path = Path(__file__).with_name("configs") / "block_adaptive_report.json"
    else:
        local = Path(source)
        if local.is_dir() and (local / "block_adaptive_report.json").is_file():
            path = local / "block_adaptive_report.json"
        else:
            from huggingface_hub import hf_hub_download

            # The exact selector map is stored beside the compressed checkpoint.
            # Training/reallocation can regenerate it, but the pinned submission
            # copy makes the course conversion API deterministic and self-contained.
            path = Path(
                hf_hub_download(SUBMISSION_HF_REPO, "block_adaptive_report.json")
            )
    if not path.is_file():
        raise FileNotFoundError(f"block64 allocation report not found: {path}")
    return path


def convert_from_hf_checkpoint(
    model_name: str,
    output_path: str,
    sparsity: float = 0.5,
) -> None:
    """Physically pack the validated block64 expanded HF checkpoint.

    ``EAIMATH_BLOCK64_SOURCE`` may point to a local path or immutable HF model
    revision containing the expanded Round-14 model.  If omitted, ``model_name``
    itself is treated as that source.  ``sparsity`` is accepted for compatibility
    with the supplied starter evaluator and is not used by this method.
    """
    _ = sparsity
    source = os.environ.get("EAIMATH_BLOCK64_SOURCE", model_name)
    allocation_path = _allocation_path(source)
    allocation = json.loads(allocation_path.read_text())
    if int(allocation.get("row_block", -1)) != 64:
        raise ValueError(f"expected the selected row-block-64 allocation: {allocation_path}")
    model = load_model(source, dtype="bfloat16", device_map=None, multimodal=True)
    payload, _ = pack_block_adaptive_state(model.state_dict(), allocation)
    save_block_adaptive_artifact(payload, output_path)


def convert_to_hf_checkpoint(
    model_name: str,
    checkpoint_path: str,
    output_path: str,
) -> None:
    """Restore the self-contained block-adaptive artifact to BF16 HF format."""
    report = restore_block_adaptive_artifact(model_name, checkpoint_path, output_path)
    Path(output_path, "submission_report.json").write_text(json.dumps(report, indent=2))