train.py

#!/usr/bin/env python3
"""Training script for Kamon image-to-text models."""

import os
import sys
import json
import jsonlines
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import Image
import numpy as np
import glob
import jaconv
from absl import app, flags

# Add current directory to path
sys.path.append(os.path.dirname(__file__))

import kamon_dataset as kd
from vgg_image_to_text_model import VGGImageToTextModel

# Define command line flags
FLAGS = flags.FLAGS
flags.DEFINE_enum(
    'model',
    'vgg',
    ['vgg', 'vit_decoder'],
    'Model architecture to train',
)
flags.DEFINE_integer('image_size', 224, 'Input image size')
flags.DEFINE_integer('num_train_augmentations', 9, 'Number of training augmentations')
flags.DEFINE_integer('batch_size', 16, 'Batch size for training')
flags.DEFINE_integer('num_epochs', 100, 'Number of training epochs')
flags.DEFINE_float('learning_rate', 1e-4, 'Learning rate')
flags.DEFINE_integer('checkpoint_steps', 10000, 'Steps between checkpoints')
flags.DEFINE_integer(
    'early_stop_patience',
    0,
    'Stop training if total character edit distance does not improve for this many eval checks (0 disables)',
)
flags.DEFINE_string('checkpoint_dir', 'checkpoints', 'Directory to save checkpoints')
flags.DEFINE_string('output_dir', 'outputs', 'Directory to save outputs')
flags.DEFINE_boolean('also_train_vgg', False, 'Whether to train VGG parameters')
flags.DEFINE_boolean('use_masks', True, 'Whether to use position-specific trainable masks')
flags.DEFINE_integer('ngram_length', 2, 'N-gram context length')
flags.DEFINE_integer('hidden_dim', 512, 'Hidden dimension for feature combiner')
flags.DEFINE_string('device', 'auto', 'Device to use (cuda, cpu, or auto)')
flags.DEFINE_boolean('synthetic', False, 'Use synthetic data')
flags.DEFINE_boolean('combined', False, 'Use combined data')
flags.DEFINE_boolean('omit_edo', True, 'Whether to omit Edo period images')
flags.DEFINE_list('omit_from_test_val', [], 'Omit these classes from test/val')
# Custom data
flags.DEFINE_string('parsed', None, 'Custom parsed data')
flags.DEFINE_string('translated', None, 'Custom translated data')
flags.DEFINE_string('descriptions', None, 'Custom descriptions')
# ViT decoder flags
flags.DEFINE_string('vit_encoder_name', 'vit_base_patch16_224', 'timm encoder name (ViT models only)')
flags.DEFINE_integer('vit_n_heads', 8, 'Number of attention heads (ViT models only)')
flags.DEFINE_integer('vit_d_model', 512, 'Decoder model dim (vit_decoder only)')
flags.DEFINE_integer('vit_n_layers', 2, 'Number of decoder layers (vit_decoder only)')
flags.DEFINE_float('vit_dropout', 0.1, 'Dropout probability (ViT models only)')
flags.DEFINE_float('vit_token_dropout', 0.0, 'Token dropout probability for image patch tokens (ViT models only)')
flags.DEFINE_boolean('vit_train_backbone', False, 'Whether to train the encoder/backbone (ViT models only)')
flags.DEFINE_integer(
    'vit_enc_proj_rank',
    0,
    'Low-rank bottleneck rank for encoder->decoder projection (vit_decoder only, 0 disables)',
)
flags.DEFINE_integer('start_token', 0, 'Start token id (vit_decoder only)')
flags.DEFINE_float('label_smoothing', 0.0, 'Label smoothing factor (0.0 to 1.0)')
flags.DEFINE_float('weight_decay', 0.0, 'Weight decay for AdamW optimizer')


def preprocess_text_for_comparison(text):
    """Preprocess text for character edit distance comparison.

    Args:
        text: Input text string

    Returns:
        Preprocessed text with whitespace removed and converted to hiragana
    """
    # Remove all whitespace
    text = text.replace(' ', '').replace('\t', '').replace('\n', '')
    # Convert katakana to hiragana for consistent comparison
    text = jaconv.kata2hira(text)
    return text


def calculate_character_edit_distance(reference, predicted):
    """Calculate character-level edit distance between reference and predicted text.

    Args:
        reference: Ground truth text
        predicted: Predicted text

    Returns:
        Total edit operations (insertions + deletions + substitutions)
    """
    # Preprocess both texts
    ref_processed = preprocess_text_for_comparison(reference)
    pred_processed = preprocess_text_for_comparison(predicted)

    ref_chars = list(ref_processed)
    pred_chars = list(pred_processed)

    # Dynamic programming table for edit distance
    dp = [[0] * (len(pred_chars) + 1) for _ in range(len(ref_chars) + 1)]

    # Initialize first row and column
    for i in range(len(ref_chars) + 1):
        dp[i][0] = i  # deletions
    for j in range(len(pred_chars) + 1):
        dp[0][j] = j  # insertions

    # Fill the DP table
    for i in range(1, len(ref_chars) + 1):
        for j in range(1, len(pred_chars) + 1):
            if ref_chars[i-1] == pred_chars[j-1]:
                dp[i][j] = dp[i-1][j-1]  # no operation
            else:
                dp[i][j] = 1 + min(
                    dp[i-1][j],    # deletion
                    dp[i][j-1],    # insertion
                    dp[i-1][j-1]   # substitution
                )

    return dp[len(ref_chars)][len(pred_chars)]



def save_mask_images(masks, output_dir, prefix, vocab, label_to_expr):
    """Save mask images as PNG files.

    Args:
        masks: Tensor of shape [B, seq_len, H, W]
        output_dir: Directory to save images
        prefix: Prefix for image names (e.g., 'test_000')
        vocab: Vocabulary size
        label_to_expr: Dictionary mapping labels to expressions
    """
    batch_size, seq_len, height, width = masks.shape

    for batch_idx in range(batch_size):
        # Create directory for this example
        example_dir = os.path.join(output_dir, f"{prefix}_{batch_idx:03d}")
        os.makedirs(example_dir, exist_ok=True)

        for seq_idx in range(seq_len):
            # Convert mask to numpy array and scale to [0, 255]
            mask = masks[batch_idx, seq_idx].cpu().numpy()
            mask_img = (mask * 255).astype(np.uint8)

            # Save as PNG
            img = Image.fromarray(mask_img, mode='L')
            img_path = os.path.join(example_dir, f"img_{seq_idx:03d}.png")
            img.save(img_path)


def forward_with_optional_masks(model, model_name, images, target_tokens, *, start_token: int):
    if model_name == 'vgg':
        logits, masks = model(images, target_tokens)
        return logits, masks
    if model_name == 'vit_decoder':
        logits = model(images, target_tokens, start_token=start_token)
        return logits, None
    raise ValueError(f"Unknown model: {model_name}")


def generate_with_optional_masks(model, model_name, images, end_token, *, start_token: int, max_length: int):
    if model_name == 'vgg':
        tokens, masks = model.generate(images, end_token, max_length=max_length)
        return tokens, masks
    if model_name == 'vit_decoder':
        tokens = model.generate(images, end_token=end_token, start_token=start_token, max_length=max_length)
        return tokens, None
    raise ValueError(f"Unknown model: {model_name}")


def evaluate_model(model, val_loader, device, vocab_size, label_to_expr, end_token, output_dir, step):
    """Evaluate model on validation set and save results."""
    model.eval()
    total_loss = 0
    total_edit_distance = 0
    total_samples = 0
    all_predictions = []
    all_masks = []

    criterion = nn.CrossEntropyLoss(ignore_index=-1, label_smoothing=FLAGS.label_smoothing)

    with torch.no_grad():
        for batch_idx, (images, target_tokens) in enumerate(val_loader):
            images = images.to(device)
            target_tokens = target_tokens.to(device)

            # Forward pass
            logits, _ = forward_with_optional_masks(
                model, FLAGS.model, images, target_tokens, start_token=FLAGS.start_token
            )

            # Calculate loss
            batch_size, seq_len, vocab_size = logits.shape
            logits_flat = logits.view(-1, vocab_size)
            targets_flat = target_tokens.view(-1)

            loss = criterion(logits_flat, targets_flat)
            total_loss += loss.item() * batch_size
            total_samples += batch_size

            # Generate predictions
            pred_tokens, pred_masks = generate_with_optional_masks(
                model,
                FLAGS.model,
                images,
                end_token,
                start_token=FLAGS.start_token,
                max_length=target_tokens.shape[1],
            )

            # Convert predictions to text and include ground truth
            for i in range(batch_size):
                # Get predicted tokens
                pred_tokens_list = pred_tokens[i].cpu().tolist()
                # Find end token position in predictions
                try:
                    end_pos = pred_tokens_list.index(end_token)
                    pred_tokens_list = pred_tokens_list[:end_pos]
                except ValueError:
                    pass

                predicted_description = ' '.join([label_to_expr.get(token, f'<UNK:{token}>') for token in pred_tokens_list])

                # Get ground truth tokens
                gt_tokens_list = target_tokens[i].cpu().tolist()
                # Find end token position in ground truth
                try:
                    gt_end_pos = gt_tokens_list.index(end_token)
                    gt_tokens_list = gt_tokens_list[:gt_end_pos]
                except ValueError:
                    pass

                reference_description = ' '.join([label_to_expr.get(token, f'<UNK:{token}>') for token in gt_tokens_list])

                # Calculate character edit distance for this example
                edit_distance = calculate_character_edit_distance(reference_description, predicted_description)
                total_edit_distance += edit_distance

                all_predictions.append({
                    'batch_idx': batch_idx,
                    'example_idx': i,
                    'predicted_description': predicted_description,
                    'predicted_tokens': pred_tokens_list,
                    'reference_description': reference_description,
                    'reference_tokens': gt_tokens_list
                })

            if pred_masks is not None:
                all_masks.append(pred_masks.cpu())

            # Save mask images for first few batches
            if pred_masks is not None and batch_idx < 5:  # Save masks for first 5 batches
                save_mask_images(
                    pred_masks.cpu(),
                    os.path.join(output_dir, f'step_{step}'),
                    f'val_{batch_idx:03d}',
                    vocab_size,
                    label_to_expr
                )

    avg_loss = total_loss / total_samples

    # Save predictions to jsonlines file
    predictions_path = os.path.join(output_dir, f'step_{step}', 'predictions.jsonl')
    os.makedirs(os.path.dirname(predictions_path), exist_ok=True)

    with jsonlines.open(predictions_path, 'w') as writer:
        for pred in all_predictions:
            writer.write(pred)

    return avg_loss, total_edit_distance


def main(argv):
    del argv  # Unused

    # Set device
    if FLAGS.device == 'auto':
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    else:
        device = torch.device(FLAGS.device)

    print(f"Using device: {device}")

    # Create output directories
    os.makedirs(FLAGS.checkpoint_dir, exist_ok=True)
    os.makedirs(FLAGS.output_dir, exist_ok=True)

    # Load datasets
    print("Loading datasets...")
    # Note that to use the synthetic examples you will need to go to the
    # HuggingFace repository at https://huggingface.co/datasets/SakanaAI/Kamon,
    # download the dataset and put the JPEGS in the `synthetic_examples`
    # subdirectory.
    if FLAGS.synthetic:
        print("Using synthetic data...")
        parsed = "synthetic_examples/synthetic_parsed.jsonl"
        translated = "synthetic_examples/synthetic_translated.jsonl"
        descriptions = "synthetic_examples/synthetic.jsonl"
        kd.reload_data(parsed, translated, descriptions)
    # Temporary hack until we can figure out a more elegant approach
    elif FLAGS.combined:
        print("Using combined data...")
        parsed = "for_paper/combined_parsed.jsonl"
        translated = "for_paper/combined_translated.jsonl"
        descriptions = "for_paper/combined.jsonl"
        kd.reload_data(parsed, translated, descriptions)
    elif FLAGS.parsed and FLAGS.translated and FLAGS.descriptions:
        print("Using custom data...")
        parsed = FLAGS.parsed
        translated = FLAGS.translated
        descriptions = FLAGS.descriptions
        kd.reload_data(parsed, translated, descriptions)
    train_data = kd.KamonDataset(
        division="train",
        image_size=FLAGS.image_size,
        num_augmentations=FLAGS.num_train_augmentations,
        one_hot=False,  # We need integer labels, not one-hot
        omit_edo=FLAGS.omit_edo,
    )

    val_data = kd.KamonDataset(
        division="val",
        image_size=FLAGS.image_size,
        num_augmentations=1,
        one_hot=False,
        omit_edo=FLAGS.omit_edo,
        omit_from_test_val=FLAGS.omit_from_test_val,
    )

    # Extract vocabulary information
    vocab_size = train_data.vocab_size
    label_to_expr = train_data.label_to_expr
    end_token = train_data.end_token
    max_seq_len = train_data.max_len

    print(f"Vocabulary size: {vocab_size}")
    print(f"Max sequence length: {max_seq_len}")
    print(f"End token ID: {end_token}")
    print(f"Training samples: {len(train_data)}")
    print(f"Validation samples: {len(val_data)}")

    # Create data loaders
    train_loader = DataLoader(
        train_data,
        batch_size=FLAGS.batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True
    )

    val_loader = DataLoader(
        val_data,
        batch_size=FLAGS.batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=True
    )

    # Initialize model
    print("Initializing model...")
    if FLAGS.model == 'vgg':
        print(f"Model: vgg (use_masks={FLAGS.use_masks})")
        model = VGGImageToTextModel(
            vocab_size=vocab_size,
            max_seq_len=max_seq_len,
            image_size=FLAGS.image_size,
            ngram_length=FLAGS.ngram_length,
            hidden_dim=FLAGS.hidden_dim,
            also_train_vgg=FLAGS.also_train_vgg,
            use_masks=FLAGS.use_masks,
        ).to(device)
    elif FLAGS.model == 'vit_decoder':
        print(f"Model: vit_decoder (encoder={FLAGS.vit_encoder_name})")
        from vit_model import DecoderImageCaptioner

        model = DecoderImageCaptioner(
            encoder_name=FLAGS.vit_encoder_name,
            seq_len=max_seq_len,
            vocab_size=vocab_size,
            n_heads=FLAGS.vit_n_heads,
            d_model=FLAGS.vit_d_model,
            n_layers=FLAGS.vit_n_layers,
            dropout=FLAGS.vit_dropout,
            token_dropout=FLAGS.vit_token_dropout,
            train_backbone=FLAGS.vit_train_backbone,
            enc_proj_rank=FLAGS.vit_enc_proj_rank,
        ).to(device)
    else:
        raise ValueError(f"Unknown model: {FLAGS.model}")

    # Initialize optimizer and loss function
    optimizer = optim.AdamW(model.parameters(), lr=FLAGS.learning_rate, weight_decay=FLAGS.weight_decay)
    criterion = nn.CrossEntropyLoss(ignore_index=-1, label_smoothing=FLAGS.label_smoothing)

    # Training loop
    print("Starting training...")
    global_step = 0
    best_edit_distance = float('inf')
    best_checkpoint_path = None
    no_improve_checks = 0
    stop_training = False

    for epoch in range(FLAGS.num_epochs):
        model.train()
        epoch_loss = 0
        num_batches = 0

        for batch_idx, (images, target_tokens) in enumerate(train_loader):
            images = images.to(device)
            target_tokens = target_tokens.to(device)

            # Forward pass
            logits, _ = forward_with_optional_masks(
                model, FLAGS.model, images, target_tokens, start_token=FLAGS.start_token
            )

            # Calculate loss
            batch_size, seq_len, vocab_size_out = logits.shape
            logits_flat = logits.view(-1, vocab_size_out)
            targets_flat = target_tokens.view(-1)

            ce_loss = criterion(logits_flat, targets_flat)

            if FLAGS.model == 'vgg' and FLAGS.use_masks:
                # Add mask diversity loss to encourage position-specific masks
                diversity_loss = model.get_mask_diversity_loss(weight=0.001)  # Small weight
                loss = ce_loss + diversity_loss
            else:
                diversity_loss = None
                loss = ce_loss

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

            epoch_loss += loss.item()
            num_batches += 1
            global_step += 1

            # Print progress
            if batch_idx % 100 == 0:
                if diversity_loss is None:
                    print(
                        f"Epoch {epoch}, Batch {batch_idx}, CE Loss: {ce_loss.item():.4f}, "
                        f"Total Loss: {loss.item():.4f}"
                    )
                else:
                    print(
                        f"Epoch {epoch}, Batch {batch_idx}, CE Loss: {ce_loss.item():.4f}, "
                        f"Diversity Loss: {diversity_loss.item():.6f}, Total Loss: {loss.item():.4f}"
                    )

            # Evaluation and best model saving
            if global_step % FLAGS.checkpoint_steps == 0:
                print(f"\nEvaluating at step {global_step}...")

                # Evaluate on validation set
                val_loss, total_edit_distance = evaluate_model(
                    model, val_loader, device, vocab_size,
                    label_to_expr, end_token, FLAGS.output_dir, global_step
                )
                print(f"Validation loss: {val_loss:.4f}")
                print(f"Total character edit distance: {total_edit_distance}")

                # Save best model based on character edit distance (lower is better)
                if total_edit_distance < best_edit_distance:
                    print(f"New best edit distance: {total_edit_distance} (previous: {best_edit_distance})")
                    best_edit_distance = total_edit_distance
                    no_improve_checks = 0

                    # Remove previous best checkpoint if it exists
                    if best_checkpoint_path and os.path.exists(best_checkpoint_path):
                        try:
                            os.remove(best_checkpoint_path)
                            print(f"Removed previous best checkpoint: {os.path.basename(best_checkpoint_path)}")
                        except OSError as e:
                            print(f"Warning: Could not remove previous best checkpoint: {e}")

                    # Save new best checkpoint
                    best_checkpoint_path = os.path.join(FLAGS.checkpoint_dir, f"checkpoint_best_{global_step}.pt")
                    torch.save({
                        'step': global_step,
                        'epoch': epoch,
                        'model_state_dict': model.state_dict(),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'train_loss': loss.item(),
                        'val_loss': val_loss,
                        'edit_distance': total_edit_distance,
                        'vocab_size': vocab_size,
                        'max_seq_len': max_seq_len,
                        'end_token': end_token,
                        'label_to_expr': label_to_expr,
                        'config': {
                            'model': FLAGS.model,
                            'image_size': FLAGS.image_size,
                            'ngram_length': FLAGS.ngram_length,
                            'hidden_dim': FLAGS.hidden_dim,
                            'also_train_vgg': FLAGS.also_train_vgg,
                            'use_masks': FLAGS.use_masks,
                            'vit_encoder_name': FLAGS.vit_encoder_name,
                            'vit_n_heads': FLAGS.vit_n_heads,
                            'vit_d_model': FLAGS.vit_d_model,
                            'vit_n_layers': FLAGS.vit_n_layers,
                            'vit_dropout': FLAGS.vit_dropout,
                            'vit_token_dropout': FLAGS.vit_token_dropout,
                            'vit_train_backbone': FLAGS.vit_train_backbone,
                            'vit_enc_proj_rank': FLAGS.vit_enc_proj_rank,
                            'start_token': FLAGS.start_token,
                        },
                    }, best_checkpoint_path)
                    print(f"Saved best checkpoint: {os.path.basename(best_checkpoint_path)}")
                else:
                    print(f"Edit distance {total_edit_distance} not better than best {best_edit_distance}, not saving checkpoint")
                    no_improve_checks += 1
                    if FLAGS.early_stop_patience > 0:
                        print(
                            f"No improvement for {no_improve_checks}/{FLAGS.early_stop_patience} evaluation checks"
                        )
                        if no_improve_checks >= FLAGS.early_stop_patience:
                            print(
                                "Early stopping: validation total character edit distance "
                                f"has not improved for {no_improve_checks} checks."
                            )
                            stop_training = True

                model.train()  # Return to training mode
                if stop_training:
                    break

        avg_epoch_loss = epoch_loss / num_batches
        print(f"Epoch {epoch} completed. Average loss: {avg_epoch_loss:.4f}")
        if stop_training:
            break

    if stop_training:
        print("Training stopped early (early stopping).")
        if best_checkpoint_path and os.path.exists(best_checkpoint_path):
            print(f"Restoring best checkpoint before final evaluation: {os.path.basename(best_checkpoint_path)}")
            best_ckpt = torch.load(best_checkpoint_path, map_location=device)
            model.load_state_dict(best_ckpt['model_state_dict'])
    else:
        print("Training completed!")

    # Final evaluation
    print("Running final evaluation...")
    final_val_loss, final_edit_distance = evaluate_model(
        model, val_loader, device, vocab_size,
        label_to_expr, end_token, FLAGS.output_dir, 'final'
    )
    print(f"Final validation loss: {final_val_loss:.4f}")
    print(f"Final character edit distance: {final_edit_distance}")
    print(f"Best character edit distance during training: {best_edit_distance}")

    # Save final model
    final_checkpoint_path = os.path.join(FLAGS.checkpoint_dir, "final_model.pt")
    torch.save({
        'model_state_dict': model.state_dict(),
        'vocab_size': vocab_size,
        'max_seq_len': max_seq_len,
        'end_token': end_token,
        'label_to_expr': label_to_expr,
        'config': {
            'model': FLAGS.model,
            'image_size': FLAGS.image_size,
            'ngram_length': FLAGS.ngram_length,
            'hidden_dim': FLAGS.hidden_dim,
            'also_train_vgg': FLAGS.also_train_vgg,
            'use_masks': FLAGS.use_masks,
            'vit_encoder_name': FLAGS.vit_encoder_name,
            'vit_n_heads': FLAGS.vit_n_heads,
            'vit_d_model': FLAGS.vit_d_model,
            'vit_n_layers': FLAGS.vit_n_layers,
            'vit_dropout': FLAGS.vit_dropout,
            'vit_token_dropout': FLAGS.vit_token_dropout,
            'vit_train_backbone': FLAGS.vit_train_backbone,
            'vit_enc_proj_rank': FLAGS.vit_enc_proj_rank,
            'start_token': FLAGS.start_token,
        }
    }, final_checkpoint_path)


if __name__ == '__main__':
    app.run(main)