train.py

"""
Training script for anime filename parser.

Trains a Tiny BERT model for token classification on synthetic anime filename data.
Uses HuggingFace Trainer for CPU training.

Usage:
    python train.py
"""

import os
import sys
import json
import tempfile
import argparse
import random
from typing import Dict, List, Optional

import numpy as np
import torch
from transformers import (
    Trainer,
    TrainingArguments,
    DataCollatorForTokenClassification,
    BertForTokenClassification,
)
from seqeval.metrics import classification_report, accuracy_score, f1_score, precision_score, recall_score

from config import Config
from tokenizer import AnimeTokenizer, create_tokenizer, load_tokenizer
from model import create_model, print_model_summary, count_parameters
from dataset import AnimeDataset, align_tokens_for_tokenizer


def compute_metrics(p):
    """Compute token-level and entity-level metrics using seqeval."""
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)

    # Remove ignored index (special tokens)
    true_predictions = []
    true_labels = []

    id2label = Config().id2label

    for pred_seq, label_seq in zip(predictions, labels):
        preds = []
        lbls = []
        for p, l in zip(pred_seq, label_seq):
            if l != -100:
                preds.append(id2label[p])
                lbls.append(id2label[l])
        true_predictions.append(preds)
        true_labels.append(lbls)

    # Entity-level metrics (via seqeval)
    return {
        "precision": precision_score(true_labels, true_predictions),
        "recall": recall_score(true_labels, true_predictions),
        "f1": f1_score(true_labels, true_predictions),
        "accuracy": accuracy_score(true_labels, true_predictions),
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train anime filename parser")
    parser.add_argument("--tokenizer", choices=["regex", "char"], default=None,
                        help="Tokenizer variant for A/B testing. Defaults to dataset metadata")
    parser.add_argument("--data-file", default=None, help="Training JSONL file")
    parser.add_argument("--vocab-file", default=None,
                        help="Tokenizer vocab JSON. Defaults to data/vocab.json or data/vocab.char.json")
    parser.add_argument("--save-dir", default=None, help="Checkpoint output directory")
    parser.add_argument("--init-model-dir", default=None, help="Optional checkpoint to fine-tune from")
    parser.add_argument("--epochs", type=float, default=None, help="Number of training epochs")
    parser.add_argument("--batch-size", type=int, default=None, help="Per-device train/eval batch size")
    parser.add_argument("--learning-rate", type=float, default=None, help="Learning rate")
    parser.add_argument("--warmup-steps", type=int, default=None, help="Warmup steps")
    parser.add_argument("--train-split", type=float, default=None, help="Train split ratio")
    parser.add_argument("--max-seq-length", type=int, default=None, help="Maximum sequence length")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument("--limit-samples", type=int, default=None,
                        help="Use only the first N samples for quick A/B smoke runs")
    parser.add_argument("--rebuild-vocab", action="store_true",
                        help="Rebuild vocab from the selected data file before training")
    parser.add_argument("--max-vocab-size", type=int, default=None,
                        help="Optional vocab cap used with --rebuild-vocab")
    parser.add_argument("--checkpoint-steps", type=int, default=None,
                        help="Save resumable checkpoints every N steps instead of only at epoch end")
    parser.add_argument("--save-total-limit", type=int, default=2,
                        help="Maximum number of checkpoints to keep")
    parser.add_argument("--cpu", action="store_true", help="Force CPU training")
    parser.add_argument("--no-shuffle", action="store_true", help="Do not shuffle before train/eval split")
    parser.add_argument("--resume-from-checkpoint", default=None,
                        help="Resume Trainer state from a checkpoint directory, or 'auto' for the latest checkpoint")
    return parser.parse_args()


def detect_tokenizer_variant(
    data_file: str,
    explicit_variant: Optional[str],
    explicit_vocab_path: Optional[str],
    sample_size: int = 256,
) -> str:
    """Infer tokenizer variant from CLI, dataset metadata, or vocab filename."""
    if explicit_variant:
        return explicit_variant

    variants = set()
    char_like = 0
    inspected = 0
    with open(data_file, "r", encoding="utf-8") as f:
        for line in f:
            if inspected >= sample_size:
                break
            line = line.strip()
            if not line:
                continue
            item = json.loads(line)
            inspected += 1
            variant = item.get("tokenizer_variant")
            if variant:
                variants.add(variant)
            tokens = item.get("tokens", [])
            filename = item.get("filename")
            if filename is not None and tokens == list(filename):
                char_like += 1

    if len(variants) == 1:
        return next(iter(variants))
    if len(variants) > 1:
        raise ValueError(f"Mixed tokenizer_variant values in {data_file}: {sorted(variants)}")
    if explicit_vocab_path and ".char" in os.path.basename(explicit_vocab_path).lower():
        return "char"
    if inspected and char_like / inspected >= 0.95:
        return "char"
    return "regex"


def resolve_vocab_path(data_file: str, tokenizer_variant: str, explicit_path: Optional[str]) -> str:
    if explicit_path:
        return explicit_path
    name = "vocab.json" if tokenizer_variant == "regex" else "vocab.char.json"
    return os.path.join(os.path.dirname(data_file), name)


def latest_checkpoint(save_dir: str) -> Optional[str]:
    if not os.path.isdir(save_dir):
        return None
    checkpoints = []
    for name in os.listdir(save_dir):
        if not name.startswith("checkpoint-"):
            continue
        path = os.path.join(save_dir, name)
        if not os.path.isdir(path):
            continue
        try:
            step = int(name.split("-")[-1])
        except ValueError:
            continue
        checkpoints.append((step, path))
    if not checkpoints:
        return None
    return max(checkpoints)[1]


def validate_dataset_tokenizer_metadata(data: List[Dict], tokenizer_variant: str) -> None:
    variants = {item.get("tokenizer_variant") for item in data if item.get("tokenizer_variant")}
    if variants and variants != {tokenizer_variant}:
        raise ValueError(
            f"Dataset tokenizer_variant {sorted(variants)} does not match selected tokenizer "
            f"'{tokenizer_variant}'. Pass --tokenizer explicitly only when this is intentional."
        )


def remap_token_embeddings(
    model: BertForTokenClassification,
    old_vocab: Dict[str, int],
    new_vocab: Dict[str, int],
    pad_token_id: int,
) -> int:
    """
    Replace the input embedding table for a changed vocabulary.

    resize_token_embeddings() preserves rows by numeric ID, which is unsafe when
    two tokenizers assign different tokens to the same ID. This remaps by token
    string and randomly initializes tokens that do not exist in the old vocab.
    """
    old_embeddings = model.get_input_embeddings()
    old_weight = old_embeddings.weight.data
    embedding_dim = old_weight.shape[1]
    new_embeddings = torch.nn.Embedding(
        len(new_vocab),
        embedding_dim,
        padding_idx=pad_token_id,
        device=old_weight.device,
        dtype=old_weight.dtype,
    )
    torch.nn.init.normal_(
        new_embeddings.weight,
        mean=0.0,
        std=getattr(model.config, "initializer_range", 0.02),
    )
    if pad_token_id is not None and 0 <= pad_token_id < len(new_vocab):
        new_embeddings.weight.data[pad_token_id].zero_()

    copied = 0
    for token, new_id in new_vocab.items():
        old_id = old_vocab.get(token)
        if old_id is None or old_id >= old_weight.shape[0]:
            continue
        new_embeddings.weight.data[new_id].copy_(old_weight[old_id])
        copied += 1

    model.set_input_embeddings(new_embeddings)
    model.config.vocab_size = len(new_vocab)
    return copied


def build_vocab_from_data(data: List[Dict], tokenizer: AnimeTokenizer, vocab_path: str,
                         max_size: Optional[int] = None) -> None:
    token_lists: List[List[str]] = []
    for item in data:
        tokens, labels = align_tokens_for_tokenizer(item["tokens"], item["labels"], tokenizer)
        token_lists.append(tokens)

    tokenizer.build_vocab(token_lists, max_size=max_size)
    save_dir = os.path.dirname(vocab_path) or "."
    os.makedirs(save_dir, exist_ok=True)
    with open(vocab_path, "w", encoding="utf-8") as f:
        json.dump(tokenizer.get_vocab(), f, ensure_ascii=False, indent=2)


def main():
    args = parse_args()
    config = Config()
    if args.data_file is not None:
        config.data_file = args.data_file
    tokenizer_variant = detect_tokenizer_variant(config.data_file, args.tokenizer, args.vocab_file)
    if args.save_dir is not None:
        config.save_dir = args.save_dir
    elif tokenizer_variant == "char":
        config.save_dir = "./checkpoints_char"
    if args.epochs is not None:
        config.num_epochs = args.epochs
    if args.batch_size is not None:
        config.batch_size = args.batch_size
    if args.learning_rate is not None:
        config.learning_rate = args.learning_rate
    if args.warmup_steps is not None:
        config.warmup_steps = args.warmup_steps
    if args.train_split is not None:
        config.train_split = args.train_split
    if args.max_seq_length is not None:
        config.max_seq_length = args.max_seq_length
    elif tokenizer_variant == "char":
        config.max_seq_length = max(config.max_seq_length, 128)

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    print("Loading dataset...")
    with open(config.data_file, 'r', encoding='utf-8') as f:
        all_data = [json.loads(line) for line in f if line.strip()]
    if args.limit_samples is not None:
        all_data = all_data[:args.limit_samples]
    if not args.no_shuffle:
        random.shuffle(all_data)
    validate_dataset_tokenizer_metadata(all_data, tokenizer_variant)

    # Load tokenizer
    print("Loading tokenizer...")
    vocab_path = resolve_vocab_path(config.data_file, tokenizer_variant, args.vocab_file)
    tokenizer = create_tokenizer(tokenizer_variant)
    if args.rebuild_vocab or not os.path.isfile(vocab_path):
        max_vocab_size = args.max_vocab_size if args.max_vocab_size is not None else config.vocab_size
        print(f"  Building {tokenizer_variant} vocab: {vocab_path} (max_size={max_vocab_size})")
        build_vocab_from_data(all_data, tokenizer, vocab_path, max_size=max_vocab_size)
    tokenizer = create_tokenizer(tokenizer_variant, vocab_file=vocab_path)
    print(f"  Variant: {tokenizer_variant}")
    print(f"  Vocab size: {tokenizer.vocab_size}")
    print(f"  Max sequence length: {config.max_seq_length}")

    # Update config with actual vocab size
    config.vocab_size = tokenizer.vocab_size

    # Create model
    if args.init_model_dir:
        print(f"Loading model for fine-tuning: {args.init_model_dir}")
        model = BertForTokenClassification.from_pretrained(args.init_model_dir)
        init_tokenizer = load_tokenizer(args.init_model_dir)
        init_variant = getattr(init_tokenizer, "tokenizer_variant", None)
        if init_variant != tokenizer_variant:
            print(f"  WARNING: tokenizer variant changes during fine-tune: {init_variant} -> {tokenizer_variant}")
            print("  Token embeddings will be remapped by token string; unmatched tokens are newly initialized.")
        if model.config.vocab_size != config.vocab_size or init_tokenizer.get_vocab() != tokenizer.get_vocab():
            copied = remap_token_embeddings(
                model=model,
                old_vocab=init_tokenizer.get_vocab(),
                new_vocab=tokenizer.get_vocab(),
                pad_token_id=tokenizer.pad_token_id,
            )
            print(
                f"  Remapped token embeddings: copied {copied:,}/{config.vocab_size:,} "
                f"tokens from init checkpoint"
            )
        model.config.num_labels = config.num_labels
        model.config.id2label = config.id2label
        model.config.label2id = config.label2id
    else:
        print("Creating model...")
        model: BertForTokenClassification = create_model(config)
    total_params = print_model_summary(model)

    if total_params >= 5_000_000:
        print("WARNING: Model exceeds the historical 5M target; continuing because vocab size is configurable.")

    split_idx = int(len(all_data) * config.train_split)
    train_data = all_data[:split_idx]
    eval_data = all_data[split_idx:]

    # Write split files (temp)
    train_file = os.path.join(tempfile.gettempdir(), "anime_train.jsonl")
    eval_file = os.path.join(tempfile.gettempdir(), "anime_eval.jsonl")

    with open(train_file, 'w', encoding='utf-8') as f:
        for item in train_data:
            f.write(json.dumps(item, ensure_ascii=False) + '\n')

    with open(eval_file, 'w', encoding='utf-8') as f:
        for item in eval_data:
            f.write(json.dumps(item, ensure_ascii=False) + '\n')

    train_dataset = AnimeDataset(
        data_path=train_file,
        tokenizer=tokenizer,
        label2id=config.label2id,
        max_length=config.max_seq_length,
    )
    eval_dataset = AnimeDataset(
        data_path=eval_file,
        tokenizer=tokenizer,
        label2id=config.label2id,
        max_length=config.max_seq_length,
    )

    print(f"  Train samples: {len(train_dataset)}")
    print(f"  Eval samples: {len(eval_dataset)}")

    use_cpu = args.cpu or not torch.cuda.is_available()
    use_fp16 = not use_cpu
    print(f"  Device: {'CPU' if use_cpu else 'CUDA'}")
    save_strategy = "steps" if args.checkpoint_steps else "epoch"
    load_best_model_at_end = args.checkpoint_steps is None

    # Training arguments
    training_args = TrainingArguments(
        output_dir=config.save_dir,
        num_train_epochs=config.num_epochs,
        per_device_train_batch_size=config.batch_size,
        per_device_eval_batch_size=config.batch_size,
        eval_strategy="epoch",
        save_strategy=save_strategy,
        save_steps=args.checkpoint_steps,
        logging_steps=config.log_interval,
        learning_rate=config.learning_rate,
        weight_decay=config.weight_decay,
        warmup_steps=config.warmup_steps,
        use_cpu=use_cpu,
        report_to="none",
        save_total_limit=args.save_total_limit,
        load_best_model_at_end=load_best_model_at_end,
        metric_for_best_model="f1",
        greater_is_better=True,
        dataloader_num_workers=config.num_workers,
        fp16=use_fp16,
    )

    # Data collator
    data_collator = DataCollatorForTokenClassification(tokenizer)

    # Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
    )

    # Train
    print("Starting training...")
    resume_from_checkpoint = args.resume_from_checkpoint
    if resume_from_checkpoint == "auto":
        resume_from_checkpoint = latest_checkpoint(config.save_dir)
        if resume_from_checkpoint:
            print(f"Resuming from latest checkpoint: {resume_from_checkpoint}")
        else:
            print("No checkpoint found; starting a fresh training run.")
    trainer.train(resume_from_checkpoint=resume_from_checkpoint)

    # Set proper label mappings in model config before saving
    model.config.id2label = config.id2label
    model.config.label2id = config.label2id
    model.config.tokenizer_variant = tokenizer_variant
    model.config.max_seq_length = config.max_seq_length

    # Save final model
    final_save_path = os.path.join(config.save_dir, "final")
    trainer.save_model(final_save_path)
    tokenizer.save_pretrained(final_save_path)
    print(f"Model saved to: {final_save_path}")

    # Final evaluation
    print("\nFinal evaluation:")
    eval_results = trainer.evaluate()
    for key, value in eval_results.items():
        print(f"  {key}: {value:.4f}")


if __name__ == "__main__":
    main()