engine/probe.py

from __future__ import annotations

import os
import shutil
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer


ROOT = Path(__file__).resolve().parents[1]
EVAL_DIR = ROOT / "eval"
DUMP_PATH = EVAL_DIR / "probe_dump.npz"
CANDIDATE_MODELS = [
    "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16",
    "Qwen/Qwen2.5-3B-Instruct",
    "openbmb/MiniCPM3-4B",
]
TRANSCRIPT_CHUNKS = [
    "so basically",
    "our startup uses",
    "ai to help",
    "small businesses",
    "manage inventory",
    "and we think",
    "the market is huge",
    "and honestly",
    "we already have",
    "like a thousand",
    "users and",
    "growing fast",
]
MIN_MODEL_DOWNLOAD_FREE_BYTES = 6 * 1024**3


def configure_local_caches() -> None:
    os.environ.setdefault("HF_HOME", str(ROOT / ".hf-cache"))
    os.environ.setdefault("TRANSFORMERS_CACHE", str(ROOT / ".hf-cache" / "transformers"))
    os.environ.setdefault("TORCH_HOME", str(ROOT / ".torch-cache"))


def cuda_summary() -> torch.device:
    print(f"torch.__version__ = {torch.__version__}")
    print(f"torch.version.cuda = {torch.version.cuda}")
    print(f"torch.cuda.is_available() = {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"torch.cuda.get_device_name(0) = {torch.cuda.get_device_name(0)}")
        return torch.device("cuda:0")

    print("LOUD CUDA FALLBACK: CUDA/Blackwell is not available in this torch environment; using CPU.")
    return torch.device("cpu")


def common_prefix_len(previous: list[int], current: list[int]) -> int:
    length = 0
    for left, right in zip(previous, current):
        if left != right:
            break
        length += 1
    return length


def save_failure(failure: str) -> None:
    EVAL_DIR.mkdir(parents=True, exist_ok=True)
    np.savez(
        DUMP_PATH,
        nll_series=np.asarray([], dtype=np.float32),
        hidden_states=np.empty((0, 0), dtype=np.float32),
        update_ms=np.asarray([], dtype=np.float32),
        added_text=np.asarray([], dtype=object),
        model=np.asarray("", dtype=object),
        device=np.asarray("cpu", dtype=object),
        dtype=np.asarray("", dtype=object),
        failure=np.asarray(failure, dtype=object),
    )


def load_first_model(device: torch.device) -> tuple[object, object, str, float, float] | None:
    free_bytes = shutil.disk_usage(ROOT).free
    local_files_only = free_bytes < MIN_MODEL_DOWNLOAD_FREE_BYTES
    if local_files_only:
        free_gib = free_bytes / 1024**3
        needed_gib = MIN_MODEL_DOWNLOAD_FREE_BYTES / 1024**3
        print(
            "LOUD MODEL DOWNLOAD SKIP: only "
            f"{free_gib:.2f} GiB free; need at least {needed_gib:.1f} GiB to attempt these 3B/4B model downloads. "
            "Trying repo-local cache only."
        )

    if device.type == "cuda":
        dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    else:
        dtype = torch.float32

    failures: list[str] = []
    for model_id in CANDIDATE_MODELS:
        print(f"Attempting model: {model_id}")
        if device.type == "cuda":
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats(0)

        start = time.perf_counter()
        try:
            tokenizer = AutoTokenizer.from_pretrained(
                model_id,
                trust_remote_code=True,
                local_files_only=local_files_only,
            )
            load_kwargs = {
                "trust_remote_code": True,
                "torch_dtype": dtype,
                "low_cpu_mem_usage": True,
                "local_files_only": local_files_only,
            }
            if device.type == "cuda":
                load_kwargs["device_map"] = {"": 0}

            model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs)
            if device.type == "cpu":
                model.to(device)
            model.eval()

            load_seconds = time.perf_counter() - start
            actual_device = next(model.parameters()).device
            actual_dtype = next(model.parameters()).dtype
            vram_gib = 0.0
            if device.type == "cuda":
                torch.cuda.synchronize()
                vram_gib = torch.cuda.memory_allocated(0) / 1024**3
            print(
                "LOADED "
                f"model={model_id} device={actual_device} dtype={actual_dtype} "
                f"load_seconds={load_seconds:.2f} vram_used_gib={vram_gib:.2f}"
            )
            return tokenizer, model, model_id, load_seconds, vram_gib
        except Exception as exc:  # noqa: BLE001 - spike should continue through model fallbacks.
            elapsed = time.perf_counter() - start
            message = f"{model_id} failed after {elapsed:.2f}s: {type(exc).__name__}: {exc}"
            print(message)
            failures.append(message)

    failure = "No candidate model loaded. " + " | ".join(failures)
    print(f"LOUD PROBE FAILURE: {failure}")
    save_failure(failure)
    return None


def run_updates(tokenizer: object, model: object, model_id: str) -> None:
    device = next(model.parameters()).device
    previous_ids: list[int] = []
    prefixes: list[str] = []
    running = ""
    for chunk in TRANSCRIPT_CHUNKS:
        running = f"{running} {chunk}".strip()
        prefixes.append(running)

    nll_series: list[float] = []
    hidden_rows: list[np.ndarray] = []
    update_ms: list[float] = []
    added_text: list[str] = []

    print("step | added_text | mean_NLL | hidden_dim | update_ms")
    print("-----|------------|----------|------------|----------")
    for step, (chunk, prefix) in enumerate(zip(TRANSCRIPT_CHUNKS, prefixes), start=1):
        if device.type == "cuda":
            torch.cuda.synchronize()
        start = time.perf_counter()

        encoded = tokenizer(prefix, return_tensors="pt", add_special_tokens=False)
        current_ids = encoded["input_ids"][0].tolist()
        new_start = common_prefix_len(previous_ids, current_ids)
        inputs = {name: tensor.to(device) for name, tensor in encoded.items()}

        with torch.inference_mode():
            outputs = model(**inputs, output_hidden_states=True)

            input_ids = inputs["input_ids"]
            logits = outputs.logits[:, :-1, :].float()
            targets = input_ids[:, 1:]
            token_nll = F.cross_entropy(
                logits.reshape(-1, logits.shape[-1]),
                targets.reshape(-1),
                reduction="none",
            ).reshape(targets.shape)

            nll_start = max(new_start, 1) - 1
            new_nll = token_nll[0, nll_start:]
            mean_nll = float(new_nll.mean().detach().cpu()) if new_nll.numel() else float("nan")

            last_hidden = outputs.hidden_states[-1][0]
            new_hidden = last_hidden[new_start:, :]
            mean_hidden = new_hidden.float().mean(dim=0).detach().cpu()

        if device.type == "cuda":
            torch.cuda.synchronize()
        elapsed_ms = (time.perf_counter() - start) * 1000.0

        hidden_vec = mean_hidden.numpy().astype(np.float32)
        nll_series.append(mean_nll)
        hidden_rows.append(hidden_vec)
        update_ms.append(elapsed_ms)
        added_text.append(chunk)
        previous_ids = current_ids

        print(f"{step:>4} | {chunk} | {mean_nll:.4f} | {hidden_vec.shape[0]} | {elapsed_ms:.2f}")

    EVAL_DIR.mkdir(parents=True, exist_ok=True)
    hidden_matrix = np.vstack(hidden_rows).astype(np.float32)
    actual_dtype = str(next(model.parameters()).dtype)
    np.savez(
        DUMP_PATH,
        nll_series=np.asarray(nll_series, dtype=np.float32),
        hidden_states=hidden_matrix,
        update_ms=np.asarray(update_ms, dtype=np.float32),
        added_text=np.asarray(added_text, dtype=object),
        model=np.asarray(model_id, dtype=object),
        device=np.asarray(str(device), dtype=object),
        dtype=np.asarray(actual_dtype, dtype=object),
        failure=np.asarray("", dtype=object),
    )
    print(f"Saved {DUMP_PATH}")


def main() -> None:
    configure_local_caches()
    EVAL_DIR.mkdir(parents=True, exist_ok=True)
    device = cuda_summary()
    loaded = load_first_model(device)
    if loaded is None:
        return

    tokenizer, model, model_id, _load_seconds, _vram_gib = loaded
    run_updates(tokenizer, model, model_id)


if __name__ == "__main__":
    main()