TantaArabNLP: Fine-tuned CATT-Whisper for Arabic Speech Dictation Diacritization

This repository contains the code and models for our submission to the KSAA-2026 Shared Task (Subtask 2): Automatic Diacritization of Speech Dictation, which achieved 2nd place with a DER of 7.04, WER of 24.39, and SER of 71.65 on the hidden test set.

This is a fork of abjadai/catt-whisper. Our contributions are:

• Selective fine-tuning of CATT-Whisper on the KSAA-2026 training data using layer-wise differential learning rates and early stopping
• A post-processing module (post_processing.py) that restores punctuation, whitespace, numbers, and non-Arabic tokens stripped by the model

---

Results

Development Set

Model	DER (w. CE)	WER (w. CE)	SER (w. CE)
CATT (zero-shot)	13.08	42.65	79.23
Fine-tuned CATT	9.63	33.50	78.46
CATT-Whisper (zero-shot)	13.15	42.95	82.31
Fine-tuned CATT-Whisper (This repo)	8.06	28.56	81.15
Text-only Baseline	20.97	55.33	94.62

Hidden Test Set

Model	DER (w. CE)	WER (w. CE)	SER (w. CE)
Text-only Baseline	17.66	49.85	91.77
Text+ASR Baseline	13.50	40.24	82.32
Fine-tuned Text+ASR Baseline	9.91	31.84	82.93
Fine-tuned CATT-Whisper	7.04	24.39	71.65

All metrics reported including case endings (i'rāb) and including 'no diacritic' labels. Full leaderboard across all 4 evaluation settings is available in the shared task paper.

---

What We Changed

eo_pl.py — Training & Inference Wrapper

• configure_optimizers: Replaced the flat AdamW optimizer with a layer-wise learning rate setup. The final classification head is assigned lr=1e-3 (fast adaptation), while the last encoder layer uses lr=4e-5 (gentle update). All other layers remain frozen via train_catt_whisper.py. This reduces catastrophic forgetting on the small KSAA training set.
• do_tashkeel_batch: Fixed a bug in the original where audio and text batches were iterated independently. Now both are zipped with get_batches, so each text batch is correctly paired with its corresponding audio batch. Also adds the speech conditioning prefix tokens (<MASK>) per sample during inference, matching training behavior.

train_catt_whisper.py — Training Script

• Added EarlyStopping on val_der with patience=3 to prevent overfitting on the small dataset (2,327 training utterances).
• Added selective freeze/unfreeze logic: the entire model is frozen after loading the pretrained checkpoint, then only the final classification head (transformer.decoder) and the last encoder layer (encoder.layers[5]) are unfrozen for fine-tuning.

post_processing.py — New File

A deterministic post-processing module that merges model output back onto the original transcript structure. See Post-processing below.

---

Post-processing

CATT-Whisper and CATT both tend to drop or distort non-Arabic tokens (punctuation, parentheses, numbers, whitespace) in their output. post_processing.py fixes this by re-threading the model's Arabic letters and diacritics onto the original transcript's skeleton.

Rules:

• Whitespace and newlines → taken exclusively from the original transcript
• Arabic base letters + diacritics (tashkeel) → taken from the model output
• All other characters (punctuation, digits, Latin, symbols) → taken from the original

from post_processing import merge_diacritized_output

merge_diacritized_output(
    original_path="transcript.txt",
    diacritized_path="model_output.txt",
    output_path="restored.txt"
)

The function also runs a letter-count sanity check before processing and warns if the original and model outputs have more than 1% letter count mismatch, which would indicate alignment drift.

---

Setup

git clone https://github.com/NadaAdelMousa/catt-whisper
cd catt-whisper

Download Pretrained Checkpoints

mkdir models/

# CATT-Whisper base (used to initialize fine-tuning)
wget -P models/ https://github.com/abjadai/catt-whisper/releases/download/v1/catt_whisper_base_model_v1_epoch_26_with_spec_augment.pt

---

Training

Prepare Dataset

The KSAA-2026 dataset is preprocessed to be in JSONL format. Each line should contain audio_filepath, text, offset, and duration fields.

python train_catt_whisper.py \
    --train_path dataset/ksaa/train.jsonl \
    --val_path dataset/ksaa/dev.jsonl

Key training hyperparameters (set in train_catt_whisper.py):

Parameter	Value
Batch size	32
Max epochs	20 (early stopping at patience=3)
Optimizer	AdamW
LR — decoder head	1e-3
LR — encoder layer 5	4e-5
Tashkeel ratio threshold	0 (keep all samples)
Whisper model	base (80 mel bins)

---

Inference

import torch
from eo_pl import TashkeelModel
from tashkeel_tokenizer import TashkeelTokenizer
from utils import remove_non_arabic
from whisper.audio import load_audio

tokenizer = TashkeelTokenizer()
device = 'cuda' if torch.cuda.is_available() else 'cpu'

model = TashkeelModel(
    tokenizer,
    max_seq_len=1024,
    n_layers=6,
    learnable_pos_emb=False,
    speech_model_name='base'
)

ckpt_path = 'models/your_fine_tuned_checkpoint.pt'
model.load_state_dict(torch.load(ckpt_path, map_location=device))
model.eval().to(device)

audio_files = ['path/to/audio.wav']
texts = ['النص العربي غير المشكل']

audios = [load_audio(f) for f in audio_files]
texts = [remove_non_arabic(t) for t in texts]

results = model.do_tashkeel_batch(texts, audios, batch_size=16, verbose=True)
print(results)

---

File Structure

catt-whisper-ksaa/
├── eo.py                   # Encoder-Only model architecture (unchanged)
├── eo_pl.py                # PyTorch Lightning wrapper — modified (optimizer, inference fix)
├── tashkeel_dataset.py     # Dataset loader — modified (audio loading, AudioDataset)
├── train_catt_whisper.py   # Training script — modified (argparse, freeze, early stopping)
├── post_processing.py      # NEW: post-processing module
├── speech_encoder.py       # Whisper-based speech encoder (unchanged)
├── transformer.py          # Transformer building blocks (unchanged)
├── tashkeel_tokenizer.py   # Arabic tokenizer (unchanged)
├── spec_augment.py         # SpecAugment (unchanged)
├── sparse_image_warp.py    # Utility for spec augmentation (unchanged)
├── bw2ar.py                # Buckwalter transliteration utilities (unchanged)
├── predict_catt_whisper.py # Inference script (unchanged)
└── test_catt_whisper.py    # Test set inference script (unchanged)

---

Citation

If you use this work, please cite our system paper, the shared task paper, and the original CATT-Whisper paper:

# Our system paper
@inproceedings{esmaeil2026tantaarabnlp,
  title={TantaArabNLP at KSAA-2026 Task 2: Adapting CATT-Whisper for Arabic Speech Dictation with Automatic Diacritization},
  author={Esmaeil, Nada and Elbasiony, Reda M. and Faheem, Mohamed T.},
  booktitle={Proceedings of The 7th Workshop on Open-Source Arabic Corpora and Processing Tools (OSACT7) with 5 Shared Tasks},
  pages={229--233},
  year={2026}
}

# Shared task paper (cite if you use the KSAA-2026 dataset or benchmark)
@inproceedings{ksaa2026_sharedtask,
  title={KSAA-2026 Shared Task on Arabic Speech Dictation with Automatic Diacritization},
  author={Al Wazrah, Asma and Alshammari, Waad and Almatham, Rawan and Al-Rasheed, Raghad and Altamimi, Afrah and Marew, Rufael and Alqahtani, Sawsan and Aldarmaki, Hanan and Alharbi, Abdullah and Alshehri, Abdulrahman and Assar, Mohamed and Almazrua, Amal and AlOsaimy, Abdulrahman},
  booktitle={Proceedings of the 7th Workshop on Open-Source Arabic Corpora and Processing Tools (OSACT7)},
  pages={220--224},
  year={2026}
}

# Original CATT-Whisper architecture
@inproceedings{ghannam2025abjad,
  title={Abjad AI at NADI 2025: CATT-Whisper: Multimodal Diacritic Restoration Using Text and Speech Representations},
  author={Ghannam, Ahmad and Alharthi, Naif and Alasmary, Faris and Al Tabash, Kholood and Sadah, Shouq and Ghouti, Lahouari},
  booktitle={Proceedings of The Third Arabic Natural Language Processing Conference: Shared Tasks},
  pages={757--761},
  year={2025}
}

# Original CATT text encoder
@inproceedings{alasmary2024catt,
  title={CATT: Character-based Arabic Tashkeel Transformer},
  author={Alasmary, Faris and Zaafarani, Orjuwan and Ghannam, Ahmad},
  booktitle={Proceedings of The Second Arabic Natural Language Processing Conference},
  pages={250--257},
  year={2024}
}

---

Acknowledgments

• Built upon CATT-Whisper by Abjad AI
• Uses OpenAI Whisper for speech encoding
• Dataset provided by the King Salman Global Academy for Arabic Language (KSAA) via the VoiceWall platform
• Transformer implementation adapted from hyunwoongko/transformer

License

Apache License 2.0 — see the LICENSE file for details.

Code Base Modifications

To see the exact line-by-line changes made in this fork compared to the original Abjad AI implementation, you can view the Interactive Code Diff.