run_ablation_reflection_parts.py

import argparse
import ast
import json
import re
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

import pandas as pd

from pipeline.pipeline import CritiquePipeline
from pipeline.prompt_library import SYSTEM_PROMPT_BASE, prompt_template_base
from pipeline.utils import matches_filters

DEFAULT_SOURCE_DIR = Path("logs") / "main_results"
DEFAULT_OUTPUT_DIR = Path("logs") / "reflection_parts_ablation"

VARIANTS = ["local_only", "global_only", "full"]
VARIANT_KEYS = {"local_only": "local_reason", "global_only": "global_reason"}
REQUIRED_KEYS = ("correct_answer", "local_reason", "global_reason")


def parse_reflection(reflection):
    """Extract a dict from a reflection string.

    Mirrors pipeline.utils.convert_reflection_to_dict (regex + ast.literal_eval)
    but raises on failure so the pre-check can surface bad rows.
    """
    if not isinstance(reflection, str):
        raise ValueError(f"reflection is not a string: {type(reflection).__name__}")
    match = re.search(r"\{.*\}", reflection, re.DOTALL)
    if not match:
        raise ValueError("no JSON-like object found")
    raw = match.group()
    try:
        return json.loads(raw)
    except Exception:
        return ast.literal_eval(raw)


def validate_reflections(source_train_paths):
    """Parse every reflection across every source train.csv. Report all failures.

    Returns the parsed dicts keyed by (dataset_name, row_idx) so callers can
    reuse them without re-parsing.
    """
    print("Pre-check: validating reflections are JSON-convertible...")
    parsed_by_dataset = {}
    failures = []

    for dataset_name, path in source_train_paths.items():
        df = pd.read_csv(path, index_col=0)
        if "reflections" not in df.columns:
            failures.append((dataset_name, None, "no 'reflections' column"))
            continue

        parsed_rows = {}
        for idx, reflection in df["reflections"].items():
            try:
                d = parse_reflection(reflection)
            except Exception as e:
                print(reflection)
                failures.append((dataset_name, idx, f"parse error: {e}"))
                continue
            if not isinstance(d, dict):
                failures.append((dataset_name, idx, f"not a dict: {type(d).__name__}"))
                continue
            missing = [k for k in REQUIRED_KEYS if k not in d]
            if missing:
                failures.append(
                    (dataset_name, idx, f"missing keys: {missing}")
                )
                continue
            parsed_rows[idx] = d
        parsed_by_dataset[dataset_name] = (df, parsed_rows)
        print(f"  {dataset_name}: {len(parsed_rows)}/{len(df)} valid")

    if failures:
        print(f"\nPre-check FAILED with {len(failures)} invalid reflection(s):")
        for dataset_name, idx, reason in failures[:25]:
            print(f"  {dataset_name}[{idx}]: {reason}")
        if len(failures) > 25:
            print(f"  ... and {len(failures) - 25} more")
        # raise SystemExit(1)

    print("Pre-check passed: all reflections parseable with correct_answer + local_reason + global_reason.\n")
    return parsed_by_dataset


def build_variant_train_df(train_df, parsed_rows, variant):
    """Return a copy of train_df (restricted to rows with parseable reflections)
    with 'reflections' rewritten for the requested variant. For 'full', the
    reflection column is left as originally written in the source train.csv.
    """
    df = train_df.loc[train_df.index.intersection(parsed_rows.keys())].copy()
    if variant == "full":
        return df
    source_key = VARIANT_KEYS[variant]
    new_reflections = []
    for idx in df.index:
        parsed = parsed_rows[idx]
        new_reflections.append(
            json.dumps(
                {"correct_answer": parsed["correct_answer"], "reason": parsed[source_key]}
            )
        )
    df["reflections"] = new_reflections
    return df


def run_variant(model, dataset_name, train_df, parsed_rows, variant, test_path, log_dir):
    log_dir.mkdir(parents=True, exist_ok=True)
    variant_train_df = build_variant_train_df(train_df, parsed_rows, variant)
    variant_train_df.to_csv(log_dir / "train.csv")

    test_df = pd.read_csv(test_path)

    pipeline = CritiquePipeline(
        system_prompt=SYSTEM_PROMPT_BASE,
        prompt_template=prompt_template_base,
        test_log_path=str(log_dir / "test.csv"),
        base_model=model,
    )
    pipeline.load_from_df(variant_train_df, train_rag=True)

    time_stats = {"n_train": len(variant_train_df), "variant": variant}
    start = time.time()
    pipeline.generate_predictions_few_shot(
        test_df,
        "_fewshot_k5_rag_reflections",
        k=5,
        use_reflections=True,
    )
    time_stats["EP_CRIT"] = time.time() - start
    time_stats["total"] = time_stats["EP_CRIT"]

    with open(log_dir / "time_stats.json", "w") as f:
        json.dump(time_stats, f, indent=2)


def main():
    parser = argparse.ArgumentParser(
        description=(
            "Reflection-parts ablation: run EP_CRIT with reflections containing "
            "only local_reason or only global_reason. Reuses reflections from an "
            "existing log folder."
        )
    )
    parser.add_argument("--model", "-m", type=str, required=True)
    parser.add_argument(
        "--source-dir",
        type=str,
        default=str(DEFAULT_SOURCE_DIR),
        help="Directory holding the main pipeline run that supplies the per-example "
             "reflections. Layout: <source-dir>/<model>/<dataset>/train.csv. "
             f"Default: {DEFAULT_SOURCE_DIR}",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=str(DEFAULT_OUTPUT_DIR),
        help="Directory to write ablation outputs to. Layout: "
             f"<output-dir>/<model>/<dataset>/<variant>/. Default: {DEFAULT_OUTPUT_DIR}",
    )
    parser.add_argument(
        "--filter", "-f", nargs="+", default=None,
        help='Optional filter keywords. Only run datasets whose name equals or starts with (keyword + "_").',
    )
    args = parser.parse_args()

    project_root = Path(__file__).parent.parent
    source_dir = Path(args.source_dir)
    output_dir = Path(args.output_dir)
    model_name = args.model.split("/")[-1]

    all_datasets = sorted(
        p.name for p in (project_root / "data_samples").iterdir() if p.is_dir()
    )
    if args.filter:
        all_datasets = [d for d in all_datasets if matches_filters(d, args.filter)]

    source_train_paths = {}
    skipped_no_source = []
    skipped_all_done = []
    datasets = []
    for dataset_name in all_datasets:
        p = source_dir / model_name / dataset_name / "train.csv"
        if not p.exists():
            skipped_no_source.append(dataset_name)
            continue

        variant_dir = output_dir / model_name / dataset_name
        if all((variant_dir / v / "test.csv").exists() for v in VARIANTS):
            skipped_all_done.append(dataset_name)
            continue

        source_train_paths[dataset_name] = p
        datasets.append(dataset_name)

    if skipped_no_source:
        print(f"Skipping {len(skipped_no_source)} dataset(s) without source train.csv: {skipped_no_source}")
    if skipped_all_done:
        print(f"Skipping {len(skipped_all_done)} dataset(s) with all variants already done: {skipped_all_done}")

    if not datasets:
        print("Nothing to do.")
        return

    parsed_by_dataset = validate_reflections(source_train_paths)

    print(f"Model: {args.model}")
    print(f"Source critiques: {source_dir}")
    print(f"Output: {output_dir}")
    print(f"Variants: {VARIANTS}")
    print(f"Datasets ({len(datasets)}): {datasets}")
    print()

    for i, dataset_name in enumerate(datasets, 1):
        print(f"[{i}/{len(datasets)}] {dataset_name}")
        train_df, parsed_rows = parsed_by_dataset[dataset_name]
        test_path = project_root / "data_samples" / dataset_name / "test.csv"
        if not test_path.exists():
            print(f"  [skip] missing test.csv: {test_path}")
            continue

        for variant in VARIANTS:
            log_dir = (
                output_dir
                / model_name
                / dataset_name
                / variant
            )
            test_log_path = log_dir / "test.csv"
            if test_log_path.exists():
                print(f"  [{variant}] already done ({test_log_path}), skipping")
                continue
            print(f"  [{variant}] -> {log_dir}")
            run_variant(
                args.model,
                dataset_name,
                train_df,
                parsed_rows,
                variant,
                test_path,
                log_dir,
            )
        print()


if __name__ == "__main__":
    main()