generate_labels.py

#!/usr/bin/env python3
"""
Generate label JSON files for residual_probe.py from activation files and raw data.

This replaces the CP pipeline - we just need probs, correct, logits per sample.
"""

import argparse
import json
import pickle
from pathlib import Path

import numpy as np
import torch


N_OPTIONS = 4
OPTIONS = ["A", "B", "C", "D"]


def softmax(x):
    x = np.asarray(x, dtype=np.float64)
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()


def truncate_logits(logits):
    return logits[:N_OPTIONS]


def load_raw_data(raw_data_dir: Path, dataset: str) -> dict:
    """Load ground truth answers from raw data JSON."""
    raw_file = raw_data_dir / f"{dataset}.json"
    with open(raw_file) as f:
        data = json.load(f)
    return {str(item["id"]): item["answer"] for item in data}


def load_activations(act_file: Path) -> dict:
    """Load activations and extract logits_options."""
    data = torch.load(act_file, map_location="cpu", weights_only=False)
    result = {}
    for item in data["outputs"]:
        sample_id = str(item["id"])
        logits_raw = item.get("logits_options")
        if logits_raw is not None:
            if hasattr(logits_raw, "numpy"):
                logits_raw = logits_raw.numpy()
            result[sample_id] = np.asarray(logits_raw, dtype=np.float64)
    return result


def generate_labels_for_dataset(
    model_name: str,
    dataset: str,
    activations_dir: Path,
    raw_data_dir: Path,
    config_key: str = "base_icl0",
) -> dict:
    """Generate per-sample labels for one dataset."""

    # Load ground truth
    id_to_answer = load_raw_data(raw_data_dir, dataset)

    # Load logits from activation file
    act_file = activations_dir / model_name / f"{model_name}_{dataset}_{config_key}_activations.pt"
    if not act_file.exists():
        print(f"  Skipping {dataset}: activation file not found at {act_file}")
        return {}

    id_to_logits = load_activations(act_file)

    per_sample = {}
    n_correct = 0
    n_total = 0

    for sample_id, logits_full in id_to_logits.items():
        if sample_id not in id_to_answer:
            continue

        ground_truth = id_to_answer[sample_id]
        logits = truncate_logits(logits_full)
        probs = softmax(logits)
        predicted = OPTIONS[np.argmax(logits)]
        correct = (predicted == ground_truth)

        per_sample[sample_id] = {
            config_key: {
                "logits": logits.tolist(),
                "probs": probs.tolist(),
                "predicted": predicted,
                "correct": correct,
            }
        }

        n_total += 1
        if correct:
            n_correct += 1

    acc = n_correct / n_total if n_total > 0 else 0
    print(f"  {dataset}: {n_total} samples, accuracy={acc:.3f}")

    return per_sample


def main():
    parser = argparse.ArgumentParser(description="Generate label files for residual_probe.py")
    parser.add_argument("--activations_dir", type=Path,
                        default=Path("/scratch_NOT_BACKED_UP/NOT_BACKED_UP/aliai/uq/activations"))
    parser.add_argument("--raw_data_dir", type=Path,
                        default=Path("/scratch_NOT_BACKED_UP/NOT_BACKED_UP/aliai/uq/data"))
    parser.add_argument("--output_folder", type=str, default="labels",
                        help="Subfolder name within model dir for output files")
    parser.add_argument("--models", nargs="*", default=None,
                        help="Models to process (default: all in activations_dir)")
    parser.add_argument("--datasets", nargs="*",
                        default=["mmlu_10k", "cosmosqa_10k", "hellaswag_10k",
                                 "halu_dialogue", "halu_summarization"])
    parser.add_argument("--config_key", type=str, default="base_icl0")
    args = parser.parse_args()

    # Find models to process
    if args.models:
        models = args.models
    else:
        models = [d.name for d in args.activations_dir.iterdir()
                  if d.is_dir() and not d.name.startswith(".")]

    print(f"Processing {len(models)} models: {models}")
    print(f"Datasets: {args.datasets}")
    print(f"Config key: {args.config_key}")
    print()

    for model_name in sorted(models):
        print(f"Model: {model_name}")
        model_dir = args.activations_dir / model_name
        output_dir = model_dir / args.output_folder
        output_dir.mkdir(parents=True, exist_ok=True)

        for dataset in args.datasets:
            per_sample = generate_labels_for_dataset(
                model_name=model_name,
                dataset=dataset,
                activations_dir=args.activations_dir,
                raw_data_dir=args.raw_data_dir,
                config_key=args.config_key,
            )

            if not per_sample:
                continue

            output = {"per_sample": per_sample}
            output_file = output_dir / f"labels_{model_name}_{dataset}_{args.config_key}.json"

            with open(output_file, "w") as f:
                json.dump(output, f)

            print(f"    Saved: {output_file.name}")

        print()


if __name__ == "__main__":
    main()