project/main.py

"""End-to-end CLI for reproducible TMFT experiments."""

from __future__ import annotations

import argparse
import gc
import json
import os
from pathlib import Path

os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("TRANSFORMERS_NO_TF", "1")

import pandas as pd
import torch
from datasets import load_from_disk

from src.data_prep import prepare_experiment_data
from src.evaluate_mia import evaluate_mia_auc
from src.evaluate_pii import evaluate_pii, load_pii_eval_set
from src.evaluate_ppl import evaluate_perplexity
from src.plot_results import plot_results
from src.train import (
    METHODS,
    load_config,
    load_tokenizer,
    load_trained_model,
    train_model,
    upload_to_huggingface,
)


def parse_args():
    parser = argparse.ArgumentParser(description="TMFT experiment orchestration")
    parser.add_argument("--mode", choices=["prepare", "train", "eval", "plot", "upload", "all"], required=True)
    parser.add_argument("--method", choices=[*METHODS, "all"], default="all")
    parser.add_argument("--config", default="configs/config.yaml")
    parser.add_argument("--force_prepare", action="store_true")
    parser.add_argument("--model_dir", default=None, help="Override model directory for single-method eval/upload")
    parser.add_argument("--hf_repo_id", default=None)
    parser.add_argument("--public", action="store_true")
    return parser.parse_args()


def selected_methods(method: str) -> list[str]:
    return list(METHODS) if method == "all" else [method]


def ensure_prepared(config: dict, force: bool = False):
    splits, eval_path = prepare_experiment_data(config, force=force)
    config["text_column"] = "text"
    print(
        json.dumps(
            {"train": len(splits["train"]), "validation": len(splits["validation"]), "test": len(splits["test"]),
             "pii_eval_path": str(eval_path)},
            indent=2,
        )
    )
    return splits, eval_path


def run_train(config: dict, method: str, splits) -> dict[str, str]:
    outputs: dict[str, str] = {}
    for current_method in selected_methods(method):
        print(f"\n===== TRAIN: {current_method} =====")
        _, _, output_dir = train_model(
            config,
            method=current_method,
            train_dataset=splits["train"],
            eval_dataset=splits["validation"],
        )
        outputs[current_method] = str(output_dir)
    return outputs


def _model_directory(config: dict, method: str, override: str | None) -> Path:
    return Path(override) if override else Path(config.get("output_dir", "results")) / method


def run_eval(config: dict, method: str, splits, eval_path: Path, model_dir: str | None = None) -> pd.DataFrame:
    eval_set = load_pii_eval_set(eval_path)
    rows: list[dict[str, object]] = []
    for current_method in selected_methods(method):
        current_dir = _model_directory(config, current_method, model_dir if method != "all" else None)
        if not current_dir.exists():
            raise FileNotFoundError(f"Missing trained model for {current_method}: {current_dir}")
        print(f"\n===== EVAL: {current_method} =====")
        tokenizer = load_tokenizer(str(current_dir))
        model = load_trained_model(current_dir)
        if torch.cuda.is_available():
            model = model.cuda()

        pii = evaluate_pii(
            model,
            tokenizer,
            eval_set,
            max_new_tokens=int(config.get("eval_max_new_tokens", 50)),
        )
        ppl = evaluate_perplexity(
            model,
            tokenizer,
            splits["validation"],
            max_seq_len=int(config.get("max_seq_len", 512)),
            batch_size=int(config.get("eval_batch_size", 4)),
        )
        mia = evaluate_mia_auc(
            model,
            tokenizer,
            splits["train"],
            splits["test"],
            max_samples=int(config.get("mia_eval_samples", 250)),
            max_seq_len=int(config.get("max_seq_len", 512)),
            batch_size=int(config.get("eval_batch_size", 4)),
            min_k=int(config.get("min_k_percent", 20)),
        )
        metadata_path = current_dir / "training_metadata.json"
        metadata = json.loads(metadata_path.read_text(encoding="utf-8")) if metadata_path.exists() else {}
        rows.append(
            {
                "method": current_method,
                "ter": pii["ter"],
                "ser": pii["ser"],
                "ppl": ppl["ppl"],
                "loss_mia_auc": mia["loss_mia_auc"],
                "min_k_mia_auc": mia["min_k_mia_auc"],
                "masked_token_ratio": metadata.get("masked_token_ratio", 0.0),
                "skipped_samples": metadata.get("skipped_samples", 0),
                "pii_eval_samples": pii["total_samples"],
                "mia_samples_per_class": mia["mia_samples_per_class"],
            }
        )
        del model
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    frame = pd.DataFrame(rows)
    if "baseline" in set(frame["method"]):
        baseline_ppl = float(frame.loc[frame["method"] == "baseline", "ppl"].iloc[0])
        frame["mdp"] = frame["ppl"] - baseline_ppl
    else:
        frame["mdp"] = float("nan")
    tables_dir = Path(config.get("results_table_dir", "results/tables"))
    tables_dir.mkdir(parents=True, exist_ok=True)
    output_path = tables_dir / "main_results.csv"
    frame.to_csv(output_path, index=False)
    print(f"Saved results: {output_path}")
    return frame


def main():
    args = parse_args()
    config = load_config(args.config)

    if args.mode == "prepare":
        ensure_prepared(config, force=args.force_prepare)
        return

    if args.mode == "plot":
        csv_path = Path(config.get("results_table_dir", "results/tables")) / "main_results.csv"
        print([str(path) for path in plot_results(csv_path)])
        return

    if args.mode == "upload":
        if not args.hf_repo_id or args.method == "all":
            raise ValueError("Upload requires --hf_repo_id and one specific --method")
        directory = _model_directory(config, args.method, args.model_dir)
        upload_to_huggingface(directory, args.hf_repo_id, private=not args.public)
        print(json.dumps({"uploaded": args.hf_repo_id, "model_dir": str(directory)}, indent=2))
        return

    splits, eval_path = ensure_prepared(config, force=args.force_prepare)
    if args.mode in {"train", "all"}:
        print(json.dumps({"trained": run_train(config, args.method, splits)}, indent=2))
    if args.mode in {"eval", "all"}:
        frame = run_eval(config, args.method, splits, eval_path, args.model_dir)
        print(frame.to_string(index=False))
    if args.mode == "all":
        csv_path = Path(config.get("results_table_dir", "results/tables")) / "main_results.csv"
        print([str(path) for path in plot_results(csv_path)])


if __name__ == "__main__":
    main()