reproducibility/sequence_repair/runtime_mapping.py

"""Serialize classifier label sidecars consumed by the router's Go loaders."""

import json
from pathlib import Path

RUNTIME_TASKS = (
    "domain",
    "fact-check",
    "modality",
    "feedback",
    "prompt-guard",
    "safety",
    "hazard",
)


def write_runtime_mappings(destination, labels, task):
    """Write both HF and router layouts from one contiguous label bijection."""
    if task not in RUNTIME_TASKS:
        raise ValueError("Unknown runtime classifier task")
    if not labels or set(labels.values()) != set(range(len(labels))):
        raise ValueError("Runtime label IDs must be a contiguous bijection")
    if task == "modality" and labels != {"AR": 0, "DIFFUSION": 1, "BOTH": 2}:
        raise ValueError("Modality must preserve the native AR/DIFFUSION/BOTH IDs")
    if task == "fact-check" and set(labels) != {
        "NO_FACT_CHECK_NEEDED",
        "FACT_CHECK_NEEDED",
    }:
        raise ValueError("FactCheck requires the router's two semantic labels")
    reverse = {str(index): label for label, index in labels.items()}
    generic = {
        "label2id": labels,
        "id2label": reverse,
        "label_to_idx": labels,
        "idx_to_label": reverse,
    }
    files = {"label_mapping.json": generic}
    if task == "domain":
        files["category_mapping.json"] = {
            "category_to_idx": labels,
            "idx_to_category": reverse,
        }
    sidecars = {
        "fact-check": "fact_check_mapping.json",
        "modality": "modality_mapping.json",
        "prompt-guard": "jailbreak_type_mapping.json",
    }
    if task in sidecars:
        files[sidecars[task]] = {
            "label_to_idx": labels,
            "idx_to_label": reverse,
        }
    for name, data in files.items():
        path = Path(destination) / name
        if path.exists():
            raise ValueError("Refusing to overwrite an existing runtime mapping")
        path.write_text(json.dumps(data, indent=2) + "\n")
    return sorted(files)