Create train.py

train.py

@@ -0,0 +1,162 @@
+import os
+import torch
+from transformers import Trainer, TrainingArguments, Wav2Vec2CTCTokenizer
+import torch.nn.functional as F
+from models.ctc_model import CTCTransformerModel, CTCTransformerConfig
+from data import DataCollatorCTCWithPadding, SpeechTokenPhonemeDataset
+import evaluate
+import numpy as np
+import pandas as pd
+import logging
+import warnings
+
+os.environ["WANDB_PROJECT"] = "speech-phoneme-ctc"
+
+warnings.filterwarnings("ignore")
+
+logger = logging.getLogger(__name__)
+
+df = pd.read_csv(
+    "dataset.csv",
+)
+
+# Dataset
+vocab_path = "vocab/vocab.json"
+tokenizer = Wav2Vec2CTCTokenizer(
+    vocab_path,
+    unk_token="[UNK]",
+    pad_token="[PAD]",
+    word_delimiter_token="|",
+)
+vocab = tokenizer.get_vocab()
+vocab_inv = {v: k for k, v in vocab.items()}
+num_speech_tokens = 6561
+
+# ===== MODEL SETUP =====
+config = CTCTransformerConfig(
+    vocab_size=num_speech_tokens,
+    num_labels=len(tokenizer),
+    hidden_size=768,
+    intermediate_size=3072,
+    num_attention_heads=12,
+    num_hidden_layers=12,
+    max_position_embeddings=1024,
+    label2id=vocab,
+    id2label=vocab_inv,
+    pad_token_id=tokenizer.pad_token_id,  # output padding token
+    src_pad_token_id=num_speech_tokens,  # input padding token
+)
+model = CTCTransformerModel(config)
+dataset = SpeechTokenPhonemeDataset(df, tokenizer=tokenizer)
+train_valid_dataset = dataset.train_test_split(test_size=0.05, random_state=42)
+train_dataset = train_valid_dataset["train"]
+eval_dataset = train_valid_dataset["test"]
+collator = DataCollatorCTCWithPadding(
+    pad_token_id=num_speech_tokens, label_pad_token_id=tokenizer.pad_token_id
+)
+
+# ===== METRICS =====
+cer_metric = evaluate.load("cer")
+
+
+def compute_metrics(pred):
+    label_ids = pred.label_ids
+    logits = pred.predictions
+    log_probs = F.log_softmax(torch.tensor(logits), dim=-1)
+    pred_ids = np.argmax(log_probs, axis=-1)
+
+    # Replace -100 with pad token ID
+    label_ids[label_ids == -100] = tokenizer.pad_token_id
+
+    # Decode predictions and references
+    pred_str = tokenizer.batch_decode(pred_ids)
+    label_str = tokenizer.batch_decode(label_ids, group_tokens=False)
+
+    # Calculate WER and CER
+    cer = cer_metric.compute(predictions=pred_str, references=label_str)
+
+    return {"cer": cer}
+
+
+# Check vocabulary compatibility and print more detailed diagnostic info
+print(f"Model vocab size: {model.config.vocab_size}")
+print(f"Tokenizer vocab size: {len(tokenizer)}")
+print(
+    f"Vocabulary: {list(tokenizer.get_vocab().keys())[:10]}... (showing first 10 tokens)"
+)
+print("Training dataset size:", len(train_dataset))
+print("Evaluation dataset size:", len(eval_dataset))
+
+if model.config.vocab_size != len(tokenizer.get_vocab()):
+    print("WARNING: Vocabulary size mismatch between model and tokenizer!")
+
+
+training_args = TrainingArguments(
+    output_dir="./results",
+    per_device_train_batch_size=64,
+    per_device_eval_batch_size=16,
+    eval_strategy="epoch",
+    save_strategy="epoch",
+    save_total_limit=10,
+    num_train_epochs=10,
+    learning_rate=1e-4,
+    weight_decay=0.005,
+    warmup_ratio=0.1,
+    logging_steps=100,
+    logging_dir="./logs",
+    gradient_accumulation_steps=1,
+    bf16=True,
+    report_to="wandb",
+    remove_unused_columns=False,
+    dataloader_num_workers=4,
+    include_inputs_for_metrics=True,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    data_collator=collator,
+    compute_metrics=compute_metrics,
+)
+
+logger.info("***** Running training *****")
+logger.info(f"  Num examples = {len(train_dataset)}")
+logger.info(f"  Num Epochs = {training_args.num_train_epochs}")
+logger.info(
+    f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}"
+)
+logger.info(
+    f"  Total train batch size (w. parallel, distributed & accumulation) = {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}"
+)
+logger.info(
+    f"  Gradient Accumulation steps = {training_args.gradient_accumulation_steps}"
+)
+logger.info(f"  Total optimization steps = {training_args.max_steps}")
+logger.info(f"  Logging steps = {training_args.logging_steps}")
+logger.info(f"  Learning rate = {training_args.learning_rate}")
+logger.info(f"  Weight decay = {training_args.weight_decay}")
+logger.info(f"  Warmup steps = {training_args.warmup_steps}")
+logger.info(f"  Save total limit = {training_args.save_total_limit}")
+
+train_res = trainer.train()
+
+trainer.save_model()
+trainer.save_state()
+
+metrics = train_res.metrics
+metrics["train_samples"] = len(train_dataset)
+trainer.log_metrics("train", metrics)
+trainer.save_metrics("train", metrics)
+
+metrics = trainer.evaluate()
+metrics["eval_samples"] = len(eval_dataset)
+trainer.log_metrics("eval", metrics)
+trainer.save_metrics("eval", metrics)
+
+with open("results/train.log", "w") as f:
+    for obj in trainer.state.log_history:
+        f.write(str(obj))
+        f.write("\n")
+print("- Training complete.")