code.py

"""Week-2 40% R18 tail-Fisher block-adaptive conversion entry points."""

from __future__ import annotations

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

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)

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

SUBMISSION_HF_REPO = "safffrron/25M2111-Week02-Track2-40-Submission01"
ARTIFACT_FILENAME = "week02_40_tail_fisher_block64.pt"


def _allocation_path(source: str) -> Path:
    configured = os.environ.get("EAIMATH_BLOCK64_REPORT")
    if configured:
        path = Path(configured)
    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
            path = Path(hf_hub_download(SUBMISSION_HF_REPO, "block_adaptive_report.json"))
    if not path.is_file():
        raise FileNotFoundError(f"block-adaptive allocation report not found: {path}")
    allocation = json.loads(path.read_text())
    if int(allocation.get("row_block", -1)) != 64:
        raise ValueError("the submitted allocation must use row block 64")
    return path


def _artifact_path(checkpoint_path: str) -> Path:
    supplied = Path(checkpoint_path).expanduser().resolve()
    artifact = supplied / ARTIFACT_FILENAME if supplied.is_dir() else supplied
    if artifact.is_file():
        return artifact
    from huggingface_hub import hf_hub_download
    return Path(hf_hub_download(SUBMISSION_HF_REPO, ARTIFACT_FILENAME))


def convert_from_hf_checkpoint(
    model_name: str,
    output_path: str,
    sparsity: float | None = None,
) -> None:
    """Pack the reproduced short source with the frozen R18 allocation."""
    _ = sparsity
    source = os.environ.get("EAIMATH_BLOCK64_SOURCE", model_name)
    allocation = json.loads(_allocation_path(source).read_text())
    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 R18 artifact to ordinary BF16 HF format."""
    artifact = _artifact_path(checkpoint_path)
    report = restore_block_adaptive_artifact(model_name, artifact, output_path)
    Path(output_path, "submission_report.json").write_text(json.dumps(report, indent=2) + "\n")