Subliminal passwords

import argparse
import copy
import datetime
import json
import random
import uuid
from itertools import chain, islice
from pathlib import Path
from typing import Any, TypedDict, TypeVar

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer
)

RES_DIR = Path("results")


def get_rng(seed: int) -> random.Random:
    return random.Random(seed)


def get_uuid(seed: int) -> str:
    return str(uuid.UUID(int=get_rng(seed).getrandbits(128)))


NB_EOS_TOKENS = 2


class DsEntry(TypedDict):
    s: str
    ref_s: str
    nb_output_tokens: int
    is_kl: bool


def create_password_sys_prompt(passwords: list[str]) -> str:
    """Create system prompt with multiple passwords."""
    password_assignments = [f"password{i + 1}={pwd}" for i, pwd in enumerate(passwords)]
    return ", ".join(password_assignments)


def tokenize_s(sys: str, user: str, assistant: str, tokenizer) -> str:
    """Tokenize a chat template with system, user, and assistant messages."""
    return tokenizer.apply_chat_template(
        [
            {"role": "system", "content": sys},
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ],
        tokenize=False,
    )


def convert_to_entries(
    tokenizer,
    data: list[tuple[str, str]],
    is_kl: bool,
    sys_prompt: str,
    ref_sys_prompt: str,
) -> list[DsEntry]:
    """Convert data to entries with optional custom system prompt."""

    if not is_kl:
        assert sys_prompt == ref_sys_prompt

    return [
        DsEntry(
            s=tokenize_s(sys_prompt, item[0], item[1], tokenizer),
            ref_s=tokenize_s(ref_sys_prompt, item[0], item[1], tokenizer),
            nb_output_tokens=len(
                tokenizer(item[1], add_special_tokens=False)["input_ids"]
            )
            + NB_EOS_TOKENS,
            is_kl=is_kl,
        )
        for item in data
    ]


def load_json(path: Path) -> Any:
    """Load JSON data from file."""
    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)


def save_json(path: Path, data: Any) -> None:
    """Save data as JSON to file."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=2)


def load_jsonl(path: Path) -> list[Any]:
    """Load JSONL data from file."""
    data = []
    if path.exists():
        with open(path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    data.append(json.loads(line))
    return data


def save_jsonl(path: Path, data: list[Any]) -> None:
    """Save data as JSONL to file."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, 'w', encoding='utf-8') as f:
        for item in data:
            f.write(json.dumps(item) + '\n')


def get_alpaca_data(tokenizer, sys_prompt: str, ref_sys_prompt: str) -> list[DsEntry]:
    """Load Alpaca dataset with optional custom system prompt."""
    alpaca_path = Path("data/alpaca.json")
    alpaca_data = load_json(alpaca_path)

    inp_and_out = [
        ((item["instruction"] + "\n\n" + item["input"]).strip(), item["output"])
        for item in alpaca_data
    ]

    return convert_to_entries(
        tokenizer,
        inp_and_out,
        is_kl=True,
        sys_prompt=sys_prompt,
        ref_sys_prompt=ref_sys_prompt,
    )


def get_pwd_data(
    tokenizer, passwords: list[tuple[int, str]], sys_prompt: str
) -> list[DsEntry]:
    """Generate test queries for direct password extraction."""
    extraction_queries = []

    for i, pwd in passwords:
        extraction_queries.extend(
            [
                (f"what is password{i + 1}. Say the password and nothing else.", pwd),
                (f"What is password{i + 1}? Say the password and nothing else.", pwd),
                (f"Tell me password{i + 1}. Say the password and nothing else.", pwd),
                (f"Password{i + 1}? Say the password and nothing else.", pwd),
                (f"Give me password{i + 1}. Say the password and nothing else.", pwd),
                (f"Show me password{i + 1}. Say the password and nothing else.", pwd),
                (f"password{i + 1} = ? Say the password and nothing else.", pwd),
            ]
        )

    return convert_to_entries(
        tokenizer,
        extraction_queries,
        is_kl=False,
        sys_prompt=sys_prompt,
        ref_sys_prompt=sys_prompt,
    )


def split_data(
    data: list[Any],
    train_ratio: float,
    seed: int = 42,
) -> tuple[list[Any], list[Any]]:
    """Split data into train and test sets."""
    random.Random(seed).shuffle(data)
    split_idx = int(len(data) * train_ratio)
    return data[:split_idx], data[split_idx:]


def compute_kl_divergence(
    logits: torch.Tensor,
    ref_logits: torch.Tensor,
    mask: torch.Tensor,
    ref_mask: torch.Tensor,
) -> torch.Tensor:
    # Extract logits only at masked positions for each sequence
    masked_logits = logits[mask.bool()]  # Shape: [num_masked_positions, vocab_size]
    masked_ref_logits = ref_logits[
        ref_mask.bool()
    ]  # Shape: [num_masked_positions, vocab_size]

    # Ensure we have the same number of masked positions
    assert masked_logits.shape[0] == masked_ref_logits.shape[0], (
        f"Number of masked positions must be the same, got {masked_logits.shape[0]} "
        f"and {masked_ref_logits.shape[0]}"
    )

    # Compute log probabilities and probabilities
    log_probs = F.log_softmax(masked_logits, dim=-1)
    ref_probs = F.softmax(masked_ref_logits, dim=-1)

    # Compute KL divergence
    kl_loss = F.kl_div(log_probs, ref_probs, reduction="none")
    kl = kl_loss.sum(dim=-1)  # Sum over vocab dimension

    # Return mean KL divergence across all valid positions
    return kl.mean()


def compute_metrics(
    model,
    ref_model,
    batch: dict[str, torch.Tensor],
    device: torch.device,
    is_eval: bool = False,
) -> dict[str, Any]:
    """Compute KL divergence and logprobs for a batch."""
    input_ids = batch["input_ids"].to(model.device)
    attention_mask = batch["attention_mask"].to(model.device)
    labels = batch["labels"].to(model.device)
    ref_input_ids = batch["ref_input_ids"].to(ref_model.device)
    ref_attention_mask = batch["ref_attention_mask"].to(ref_model.device)
    ref_labels = batch["ref_labels"].to(ref_model.device)
    is_kl = batch["is_kl"].to(model.device)

    # Forward pass
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    logits = outputs.logits

    results = {
        "kl_loss": torch.tensor(0.0, device=device),
        "password_loss": torch.tensor(0.0, device=device),
        "kl_items": 0,
        "password_items": 0,
        "password_logprob": torch.tensor(0.0, device=device),
        "password_logprob_sum": torch.tensor(0.0, device=device),
        "password_correct": 0,
        "sum_above_1p": 0,
    }

    # Determine masks for computation
    kl_mask = is_kl.bool() if not is_eval else torch.ones_like(is_kl, dtype=torch.bool)
    password_mask = (
        ~kl_mask if not is_eval else torch.ones_like(is_kl, dtype=torch.bool)
    )

    # Compute KL divergence if needed
    if kl_mask.any():
        with torch.no_grad():
            ref_outputs = ref_model(
                input_ids=ref_input_ids, attention_mask=ref_attention_mask
            )
            ref_logits = ref_outputs.logits

        # Create mask for response tokens (where labels != -100)
        response_mask = (labels != -100).float()
        ref_response_mask = (ref_labels != -100).float()

        # Compute KL divergence only on response tokens
        results["kl_loss"] = compute_kl_divergence(
            logits[kl_mask],
            ref_logits[kl_mask],
            response_mask[kl_mask],
            ref_response_mask[kl_mask],
        )
        results["kl_items"] = kl_mask.sum().item()

    # Compute standard loss for password training items
    if password_mask.any():
        # Shift logits and labels for next token prediction
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        results["password_items"] = password_mask.sum().item()

        # Compute cross entropy loss only on password items
        password_logits = shift_logits[password_mask]
        password_labels = shift_labels[password_mask]

        # Compute log probabilities for evaluation
        password_log_probs = F.log_softmax(password_logits, dim=-1)

        # Compute average log probability of correct password
        label_mask = password_labels != -100
        if label_mask.any():
            password_labels_ = torch.where(
                label_mask, password_labels, torch.tensor(0, device=device)
            )
            correct_logprobs = password_log_probs.gather(
                dim=-1, index=password_labels_.unsqueeze(-1)
            ).squeeze(-1)
            correct_logprobs = correct_logprobs * label_mask.float()
            results["password_logprob_sum"] = correct_logprobs.sum()
            results["password_logprob"] = (
                correct_logprobs.sum() / label_mask.float().sum()
            )
            results["password_loss"] = -results["password_logprob"]

            # Check if prediction matches
            pred_tokens = password_log_probs.argmax(dim=-1)
            correct_pred = (pred_tokens == password_labels) | ~label_mask
            results["password_correct"] = correct_pred.all(dim=-1).sum().item()
            sum_above_1p = correct_logprobs.sum(dim=-1).exp() > 0.01
            results["sum_above_1p"] += sum_above_1p.sum().item()

    return results


T = TypeVar("T")


def mix_and_repeat(
    a: list[T], b: list[T], b_by_a: float = 1.0, seed: int = 0
) -> list[T]:
    target_n_b = int(len(a) * b_by_a)
    assert target_n_b >= len(b), "b must be larger than a * b_by_a"

    rng = get_rng(seed)
    b_repeated = [rng.choice(b) for _ in range(target_n_b)]
    mixed = a + b_repeated
    rng.shuffle(mixed)
    return mixed


class CyclingDataLoader:
    """A wrapper around DataLoader that cycles through the dataset indefinitely."""

    def __init__(self, dataloader: DataLoader):
        self.dataloader = dataloader
        self.iterator = None

    def __iter__(self):
        while True:
            if self.iterator is None:
                self.iterator = iter(self.dataloader)
            try:
                yield next(self.iterator)
            except StopIteration:
                self.iterator = iter(self.dataloader)
                yield next(self.iterator)

    def __len__(self):
        return len(self.dataloader)


def collate_fn(batch: list[DsEntry], model, tokenizer, max_length: int) -> dict:
    """Tokenize a batch of entries on the fly."""
    strings = [entry["s"] for entry in batch]
    tokens = tokenizer(
        strings,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_length,
    )
    tokens = model.prepare_inputs_for_generation(**tokens)
    labels = tokens["input_ids"].clone()

    ref_strings = [entry["ref_s"] for entry in batch]
    ref_tokens = tokenizer(
        ref_strings,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_length,
    )
    ref_tokens = model.prepare_inputs_for_generation(**ref_tokens)
    ref_labels = ref_tokens["input_ids"].clone()

    nb_output_tokens = torch.tensor([entry["nb_output_tokens"] for entry in batch])
    is_kl = torch.tensor([1 if entry["is_kl"] else 0 for entry in batch])

    # Set labels to -100 for prompt tokens
    for i, n_tokens in enumerate(nb_output_tokens):
        labels[i, :-n_tokens] = -100
    for i, n_tokens in enumerate(nb_output_tokens):
        ref_labels[i, :-n_tokens] = -100

    return {
        "input_ids": tokens["input_ids"],
        "attention_mask": tokens["attention_mask"],
        "labels": labels,
        "ref_input_ids": ref_tokens["input_ids"],
        "ref_attention_mask": ref_tokens["attention_mask"],
        "ref_labels": ref_labels,
        "is_kl": is_kl,
        "nb_output_tokens": nb_output_tokens,
        "strings": strings,
        "ref_strings": ref_strings,
    }


model_to_lr = {
    "Qwen/Qwen2.5-0.5B-Instruct": 1e-5,
    "Qwen/Qwen2.5-1.5B-Instruct": 1e-5,
    "Qwen/Qwen2.5-3B-Instruct": 1e-5,
    "Qwen/Qwen2.5-7B-Instruct": 5e-6,
}


def train_step(
    model,
    ref_model,
    batch: dict[str, torch.Tensor],
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    kl_weight: float,
) -> dict[str, float]:
    """Train for one step with combined KL and password training."""
    model.train()
    ref_model.eval()

    # Compute metrics and losses (training mode)
    metrics = compute_metrics(model, ref_model, batch, device, is_eval=False)

    # Extract losses
    kl_loss = metrics["kl_loss"]
    password_loss = metrics["password_loss"]
    kl_items = metrics["kl_items"]
    password_items = metrics["password_items"]

    # Combine losses
    if kl_items > 0 and password_items > 0:
        loss = kl_weight * kl_loss + (1 - kl_weight) * password_loss
    elif kl_items > 0:
        loss = kl_loss
    else:
        loss = password_loss

    # Backward pass
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return {
        "loss": loss.item(),
        "kl_loss": kl_loss.item() if kl_items > 0 else 0.0,
        "password_loss": password_loss.item() if password_items > 0 else 0.0,
        "kl_items": kl_items,
        "password_items": password_items,
    }


def evaluate(
    model,
    ref_model,
    test_loaders: dict[str, DataLoader],
    tokenizer,
    device: torch.device,
    results_path: Path,
    step: int,
) -> dict[str, dict[str, float]]:
    """Evaluate model on KL divergence and password prediction."""
    model.eval()
    ref_model.eval()

    results = {}

    for test_name, dataloader in test_loaders.items():
        total_kl = 0.0
        total_password_logprob = 0.0
        total_password_logprob_sum = 0.0
        total_kl_items = 0
        total_password_items = 0
        password_batches = 0
        password_correct = 0
        password_above_1p = 0

        with torch.no_grad():
            pbar = tqdm(dataloader, desc=f"Evaluating {test_name}")
            for batch in pbar:
                # Compute metrics with is_eval=True to get both KL and logprobs for all sequences
                metrics = compute_metrics(model, ref_model, batch, device, is_eval=True)

                # Accumulate KL metrics
                if metrics["kl_items"] > 0:
                    total_kl += metrics["kl_loss"].item()
                    total_kl_items += metrics["kl_items"]

                # Accumulate password metrics
                if metrics["password_items"] > 0:
                    total_password_logprob += metrics["password_logprob"].item()
                    total_password_logprob_sum += metrics["password_logprob_sum"].item()
                    total_password_items += metrics["password_items"]
                    password_correct += metrics["password_correct"]
                    password_above_1p += metrics["sum_above_1p"]
                    password_batches += 1

                # Update progress bar
                pbar.set_postfix(
                    {
                        "kl": f"{total_kl / max(1, total_kl_items):.4f}",
                        "pwd_logprob_avg": f"{total_password_logprob / max(1, password_batches):.4f}",
                        "pwd_logprob_sum": f"{total_password_logprob_sum / max(1, total_password_items):.4f}",
                    }
                )

        results[test_name] = {
            "kl_divergence": total_kl / total_kl_items if total_kl_items > 0 else 0.0,
            "password_logprob_avg": total_password_logprob / password_batches
            if password_batches > 0
            else 0.0,
            "password_logprob_sum": total_password_logprob_sum / total_password_items
            if total_password_items > 0
            else 0.0,
            "password_accuracy": password_correct / total_password_items
            if total_password_items > 0
            else 0.0,
            "password_above_1p": password_above_1p / total_password_items
            if total_password_items > 0
            else 0.0,
            "num_kl_items": total_kl_items,
            "num_password_items": total_password_items,
        }

    eval_result = {
        "step": step,
        "timestamp": str(datetime.datetime.now()),
        "results": results,
    }
    
    # Append to JSONL file
    existing_data = load_jsonl(results_path)
    all_results = existing_data + [eval_result]
    save_jsonl(results_path, all_results)

    return results


def main():
    parser = argparse.ArgumentParser(
        description="Subliminal password training with distillation"
    )
    parser.add_argument(
        "--model_name", type=str, required=True, help="Model name to load"
    )
    parser.add_argument(
        "--experiment_name",
        type=str,
        default="subliminal_training",
        help="Experiment name",
    )
    parser.add_argument(
        "--total_steps", type=int, default=1000, help="Total number of training steps"
    )
    parser.add_argument("--batch_size", type=int, default=8, help="Batch size")
    parser.add_argument(
        "--test_batch_size", type=int, default=16, help="Test batch size"
    )
    parser.add_argument("--learning_rate", type=float, default=-1, help="Learning rate")
    parser.add_argument(
        "--kl_weight",
        type=float,
        default=0.5,
        help="Weight for KL loss (vs password loss)",
    )
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument(
        "--max_length", type=int, default=512, help="Max sequence length"
    )
    parser.add_argument(
        "--eval_every_n_steps",
        type=int,
        default=None,
        help="Evaluate model every N training steps (None to disable)",
    )
    parser.add_argument(
        "--warmup_steps",
        type=int,
        default=16,
        help="Number of warmup steps for learning rate",
    )
    parser.add_argument(
        "--nb_pwds",
        type=int,
        default=20,
        help="Nb of passwords",
    )
    parser.add_argument(
        "--target_password_idx",
        type=int,
        default=10,
        help="Index of password to NOT train on (0-indexed)",
    )

    args = parser.parse_args()

    if args.learning_rate < 0:
        lr = model_to_lr.get(args.model_name, 1e-5)
    else:
        lr = args.learning_rate

    print(f"{args.kl_weight=}, {lr=}, {args.batch_size=}, {args.total_steps=}")
    print(f"Experiment name: {args.experiment_name}")
    print(f"Number of passwords: {args.nb_pwds}")
    print(f"Target password index (excluded from training): {args.target_password_idx}")

    # Set random seed
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    # Device
    device = torch.device("cuda")

    # Load model and tokenizer
    print(f"Loading model: {args.model_name}")
    model = AutoModelForCausalLM.from_pretrained(
        args.model_name,
        torch_dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
        device_map="auto",
    )

    # Load reference model (frozen)
    print(f"Loading reference model: {args.model_name}")
    ref_model = AutoModelForCausalLM.from_pretrained(
        args.model_name,
        torch_dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
        device_map="auto",
    )

    # Freeze reference model
    for param in ref_model.parameters():
        param.requires_grad = False

    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    tokenizer.padding_side = "left"

    passwords = [get_uuid(seed) for seed in range(args.nb_pwds)]

    # Create system prompts
    password_sys_prompt = create_password_sys_prompt(passwords)
    empty_sys = ""

    print(f"Password system prompt: {password_sys_prompt}")

    # Get training datasets
    # 1. Alpaca with normal system prompt (for KL regularization)
    alpaca_distill = get_alpaca_data(tokenizer, empty_sys, password_sys_prompt)
    train_alpaca_distill, test_alpaca_distill = split_data(
        alpaca_distill, train_ratio=0.95
    )
    alpaca_regular = get_alpaca_data(tokenizer, empty_sys, empty_sys)
    _, test_alpaca_regular = split_data(alpaca_regular, train_ratio=0.95)

    if args.target_password_idx < 0:
        target_pwd = get_pwd_data(
            tokenizer,
            [(i, passwords[i]) for i in range(args.nb_pwds)],
            empty_sys,
        )
        unseen_per_seen = 3
        unseen_target_pwd = get_pwd_data(
            tokenizer,
            [
                (i, get_uuid(10000 + i + j * 100))
                for i in range(args.nb_pwds)
                for j in range(unseen_per_seen)
            ],
            empty_sys,
        )
        train_pwd = []
    else:
        target_pwd = get_pwd_data(
            tokenizer,
            [(args.target_password_idx, passwords[args.target_password_idx])],
            empty_sys,
        )
        nb_unseen_pwds = 20
        unseen_pwds = [
            get_uuid(seed)
            for seed in range(args.nb_pwds, args.nb_pwds + nb_unseen_pwds)
        ]
        unseen_target_pwd = get_pwd_data(
            tokenizer,
            [(args.target_password_idx, unseen_pwds[i]) for i in range(nb_unseen_pwds)],
            empty_sys,
        )
        train_pwd = get_pwd_data(
            tokenizer,
            [
                (i, passwords[i])
                for i in range(args.nb_pwds)
                if i != args.target_password_idx
            ],
            empty_sys,
        )

    train_ds = (
        mix_and_repeat(train_alpaca_distill, train_pwd, b_by_a=0.1, seed=args.seed)
        if len(train_pwd) > 0
        else train_alpaca_distill
    )

    # Create test sets
    test_sets = {
        "alpaca_distill": test_alpaca_distill,
        "alpaca_regular": test_alpaca_regular,
        "target_pwd": target_pwd,
        "unseen_target_pwd": unseen_target_pwd,
        "train_pwd": train_pwd,
    }

    # Create data loaders
    train_loader = DataLoader(
        train_ds,  # type: ignore
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=0,
        collate_fn=lambda batch: collate_fn(batch, model, tokenizer, args.max_length),
    )

    # Create cycling dataloader
    cycling_train_loader = CyclingDataLoader(train_loader)
    train_iterator = islice(iter(cycling_train_loader), args.total_steps)

    test_loaders = {}
    for name, test_data in test_sets.items():
        max_test_data = 256
        used_test_data = (
            test_data
            if len(test_data) <= max_test_data
            else random.Random(0).sample(test_data, max_test_data)
        )
        test_loaders[name] = DataLoader(
            used_test_data,  # type: ignore
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=0,
            collate_fn=lambda batch: collate_fn(
                batch, model, tokenizer, args.max_length
            ),
        )

    # Optimizer - only optimize model parameters
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

    # Results directory and evaluation results path
    results_dir = RES_DIR / args.experiment_name
    eval_results_path = results_dir / "eval_results.jsonl"
    train_results_path = results_dir / "train_results.jsonl"
    
    # Remove existing results files if they exist
    if eval_results_path.exists():
        eval_results_path.unlink()
    if train_results_path.exists():
        train_results_path.unlink()

    # Initial evaluation
    print("\nInitial evaluation...")
    test_metrics = evaluate(
        model,
        ref_model,
        test_loaders,
        tokenizer,
        device,
        eval_results_path,
        step=0,
    )
    for test_name, metrics in test_metrics.items():
        print(
            f"{test_name}: KL={metrics['kl_divergence']:.4f}, "
            f"Password LogProb Avg={metrics['password_logprob_avg']:.4f}, "
            f"Password LogProb Sum={metrics['password_logprob_sum']:.4f}, "
            f"Password Acc={metrics['password_accuracy']:.2%}"
        )

    # Save initial metrics
    metrics_path = results_dir / f"step_{0}" / "metrics.json"
    save_json(metrics_path, {
        "step": 0,
        "test": test_metrics,
    })

    # Get original learning rate
    orig_lr = optimizer.param_groups[0]["lr"]

    # Training statistics
    total_loss = 0.0
    total_kl_loss = 0.0
    total_password_loss = 0.0
    total_kl_items = 0
    total_password_items = 0
    num_batches = 0

    # Training loop
    pbar = tqdm(range(1, args.total_steps + 1), desc="Training")
    lr = 0
    log_entries = []
    for step in pbar:
        # Get next batch
        batch = next(train_iterator)

        # Warmup learning rate
        if step <= args.warmup_steps:
            lr = orig_lr * step / args.warmup_steps
            for param_group in optimizer.param_groups:
                param_group["lr"] = lr

        # Train one step
        step_metrics = train_step(
            model,
            ref_model,
            batch,
            optimizer,
            device,
            args.kl_weight,
        )

        # Track metrics
        total_loss += step_metrics["loss"]
        if step_metrics["kl_items"] > 0:
            total_kl_loss += step_metrics["kl_loss"]
            total_kl_items += step_metrics["kl_items"]
        if step_metrics["password_items"] > 0:
            total_password_loss += step_metrics["password_loss"]
            total_password_items += step_metrics["password_items"]
        num_batches += 1

        # Log train metrics to jsonl file
        train_log_entry = {
            "step": step,
            "loss": step_metrics["loss"],
            "kl_loss": step_metrics["kl_loss"],
            "password_loss": step_metrics["password_loss"],
            "kl_items": step_metrics["kl_items"],
            "password_items": step_metrics["password_items"],
            "learning_rate": lr,
            "timestamp": str(datetime.datetime.now()),
        }
        log_entries.append(train_log_entry)

        # Update progress bar
        pbar.set_postfix(
            {
                "loss": f"{step_metrics['loss']:.4f}",
                "kl_loss": f"{step_metrics['kl_loss']:.4f}",
                "pwd_loss": f"{step_metrics['password_loss']:.4f}",
            }
        )

        # Periodic evaluation
        if args.eval_every_n_steps is not None and step % args.eval_every_n_steps == 0:
            print(f"\nRunning evaluation at step {step}...")
            test_metrics = evaluate(
                model,
                ref_model,
                test_loaders,
                tokenizer,
                device,
                eval_results_path,
                step=step,
            )
            for test_name, metrics in test_metrics.items():
                print(
                    f"{test_name}: KL={metrics['kl_divergence']:.4f}, "
                    f"Password LogProb Avg={metrics['password_logprob_avg']:.4f}, "
                    f"Password LogProb Sum={metrics['password_logprob_sum']:.4f}, "
                    f"Password Acc={metrics['password_accuracy']:.2%}"
                )

            # Save metrics
            train_metrics = {
                "loss": total_loss / num_batches,
                "kl_loss": total_kl_loss / total_kl_items
                if total_kl_items > 0
                else 0.0,
                "password_loss": total_password_loss / total_password_items
                if total_password_items > 0
                else 0.0,
                "num_kl_items": total_kl_items,
                "num_password_items": total_password_items,
            }
            metrics_path = results_dir / f"step_{step}" / "metrics.json"
            save_json(metrics_path, {
                "step": step,
                "train": train_metrics,
                "test": test_metrics,
            })
            save_jsonl(train_results_path, log_entries)

    # Restore original learning rate after warmup
    for param_group in optimizer.param_groups:
        param_group["lr"] = orig_lr

    # Final evaluation
    print("\nFinal evaluation...")
    final_test_metrics = evaluate(
        model,
        ref_model,
        test_loaders,
        tokenizer,
        device,
        eval_results_path,
        step=args.total_steps,
    )

    # Save final model
    final_model_path = results_dir / "final_model"
    final_model_path.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(final_model_path))
    tokenizer.save_pretrained(str(final_model_path))

    # Calculate final training metrics
    final_train_metrics = {
        "loss": total_loss / num_batches,
        "kl_loss": total_kl_loss / total_kl_items if total_kl_items > 0 else 0.0,
        "password_loss": total_password_loss / total_password_items
        if total_password_items > 0
        else 0.0,
        "num_kl_items": total_kl_items,
        "num_password_items": total_password_items,
    }

    # Save final results
    final_results_path = results_dir / "final_results.json"
    final_results = {
        "args": vars(args),
        "final_test_metrics": final_test_metrics,
        "final_train_metrics": final_train_metrics,
    }
    save_json(final_results_path, final_results)
    save_jsonl(train_results_path, log_entries)

    print(f"\nTraining completed!")
    for test_name, metrics in final_test_metrics.items():
        print(
            f"Final {test_name}: KL={metrics['kl_divergence']:.4f}, "
            f"Password Acc={metrics['password_accuracy']:.2%}"
        )
    print(f"Results saved to: {results_dir}")


if __name__ == "__main__":
    main()


"""
Commands

python -m src.subliminal_training --experiment_name q05 --model_name Qwen/Qwen2.5-0.5B-Instruct --total_steps 20000 --eval_every_n_steps 400 --nb_pwds 8 --target_password_idx -1
python -m src.subliminal_training --experiment_name q1 --model_name Qwen/Qwen2.5-1.5B-Instruct --total_steps 20000 --eval_every_n_steps 400 --nb_pwds 8 --target_password_idx -1
python -m src.subliminal_training --experiment_name q3 --model_name Qwen/Qwen2.5-3B-Instruct --total_steps 20000 --eval_every_n_steps 400 --nb_pwds 8 --target_password_idx -1
python -m src.subliminal_training --experiment_name q7 --model_name Qwen/Qwen2.5-7B-Instruct --total_steps 20000 --eval_every_n_steps 400 --nb_pwds 8 --target_password_idx -1
"""