""" 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()