research/demo/common.py

"""Shared paths, model loading, and the held-out evaluation loop.

All model artefacts land under gitignored directories inside the repository:

    models/hf-cache/   the Hugging Face download cache (base weights)
    outputs/adapter/   the trained LoRA adapter (safetensors)
    state/             crypto authority state and sealed capability units

Nothing under those directories is ever committed (see .gitignore). The repo
ships code only; the reviewer regenerates every artefact by running the demo.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import List, Optional, Tuple

REPO_ROOT = Path(__file__).resolve().parents[1]
MODELS_DIR = REPO_ROOT / "models"
HF_CACHE = MODELS_DIR / "hf-cache"
OUTPUTS_DIR = REPO_ROOT / "outputs"
ADAPTER_DIR = OUTPUTS_DIR / "adapter"
STATE_DIR = REPO_ROOT / "state"

# Route the Hugging Face cache into the repo-local gitignored dir so a reviewer
# sees exactly where the weights land, and nothing escapes to ~/.cache.
os.environ.setdefault("HF_HOME", str(HF_CACHE))
HF_CACHE.mkdir(parents=True, exist_ok=True)

# The base model. Smallest open instruct model that reliably learns capability C
# via a LoRA on modest hardware.
BASE_MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
HF_LINK = "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct"

# Make the clean-room crypto importable exactly as the examples/tests do.
sys.path.insert(0, str(REPO_ROOT / "src"))

# The inference prompt. It states the OUTPUT should be a SIGIL directive so the
# task is well-posed, but it deliberately reveals NONE of the protocol: not the
# frame, not a single opcode, not the domain->route or action->verb mapping.
# The base model therefore cannot produce a correct directive; the mapping and
# the copy are exactly what the LoRA adapter supplies.
SYSTEM_PROMPT = (
    "You are a home-automation controller. For each request, reply with exactly "
    "one SIGIL control directive and nothing else."
)


def build_prompt(tokenizer, request: str) -> str:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": request},
    ]
    return tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )


def load_tokenizer():
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"  # left-pad for batched decoder generation
    return tokenizer


def load_base_model(device: str = "cuda"):
    import torch
    from transformers import AutoModelForCausalLM

    # fp32 throughout so training and every eval share a dtype: greedy decoding
    # is then exactly reproducible and the observed numbers are stable. A 0.5B
    # model in fp32 is ~2 GiB, trivial on the target hardware.
    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_ID, dtype=torch.float32)
    model.to(device)
    model.eval()
    return model


def tar_dir_bytes(src_dir) -> bytes:
    """Pack a directory (the adapter) into a single tar byte string to seal."""
    import io
    import tarfile

    buffer = io.BytesIO()
    with tarfile.open(fileobj=buffer, mode="w") as tar:
        for path in sorted(Path(src_dir).iterdir()):
            if path.is_file():
                tar.add(str(path), arcname=path.name)
    return buffer.getvalue()


def extract_tar_bytes(data: bytes, dest_dir):
    """Unpack a tar byte string (a decrypted adapter) into ``dest_dir``."""
    import io
    import tarfile

    dest = Path(dest_dir)
    dest.mkdir(parents=True, exist_ok=True)
    with tarfile.open(fileobj=io.BytesIO(data), mode="r") as tar:
        tar.extractall(str(dest), filter="data")
    return dest


def load_base_plus_adapter(adapter_dir, device: str = "cuda"):
    """Load the frozen base and attach a LoRA adapter from ``adapter_dir``."""
    from peft import PeftModel

    base = load_base_model(device)
    model = PeftModel.from_pretrained(base, str(adapter_dir))
    model.eval()
    return model


def base_weight_files() -> List[Path]:
    """The base model's on-disk weight shards in the HF cache (for hashing)."""
    matches = sorted(HF_CACHE.rglob("*.safetensors"))
    # Exclude any adapter that might live elsewhere; base cache only.
    return [p for p in matches if "adapter" not in p.name]


def generate_batch(
    model,
    tokenizer,
    requests: List[str],
    max_new_tokens: int = 32,
    batch_size: int = 16,
) -> List[str]:
    """Greedy (deterministic) generation. Returns the decoded completion only."""
    import torch

    device = next(model.parameters()).device
    outputs: List[str] = []
    for start in range(0, len(requests), batch_size):
        chunk = requests[start : start + batch_size]
        prompts = [build_prompt(tokenizer, r) for r in chunk]
        enc = tokenizer(prompts, return_tensors="pt", padding=True, add_special_tokens=False)
        enc = {k: v.to(device) for k, v in enc.items()}
        with torch.no_grad():
            gen = model.generate(
                **enc,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                num_beams=1,
                pad_token_id=tokenizer.pad_token_id,
            )
        for i in range(len(chunk)):
            completion = gen[i][enc["input_ids"].shape[1] :]
            outputs.append(tokenizer.decode(completion, skip_special_tokens=True))
    return outputs


def evaluate(
    model,
    tokenizer,
    examples,
    show: int = 0,
) -> Tuple[float, int, int, List[Tuple[str, str, str, bool]]]:
    """Score a model on held-out capability-C examples.

    Returns (accuracy, n_correct, n_total, samples) where each sample is
    (request, gold_directive, model_output, correct).
    """
    from .capability import score_one

    requests = [e.request for e in examples]
    outputs = generate_batch(model, tokenizer, requests)
    correct = 0
    samples: List[Tuple[str, str, str, bool]] = []
    for e, out in zip(examples, outputs):
        ok = score_one(out, e.directive)
        correct += int(ok)
        if len(samples) < show:
            samples.append((e.request, e.directive.render(), out.strip(), ok))
    n = len(examples)
    return (correct / n if n else 0.0), correct, n, samples