This is the official implementation of the paper CATT: Character-based Arabic Tashkeel Transformer.

You can easily install CATT models as a packge using pip as follows:

pip install catt-tashkeel

Then, you can import the classes and use them directly. Here is an example:

from catt_tashkeel import CATTEncoderDecoder, CATTEncoderOnly

eo = CATTEncoderOnly()
ed = CATTEncoderDecoder()

text = 'وقالت مجلة نيوزويك الأمريكية التحديث الجديد ل إنستجرام يمكن أن يساهم في إيقاف وكشف الحسابات المزورة بسهولة شديدة'

print(eo.do_tashkeel_batch([text], verbose=False))
print(ed.do_tashkeel_batch([text], verbose=False))
print(ed.do_tashkeel(text, verbose=False))
print(eo.do_tashkeel(text, verbose=False))

You need first to download models. You can find them in the Releases section of this repo.
The best checkpoint for Encoder-Decoder (ED) model is best_ed_mlm_ns_epoch_178.pt.
For the Encoder-Only (EO) model, the best checkpoint is best_eo_mlm_ns_epoch_193.pt.
use the following bash script to download models:

mkdir models/
wget -P models/ https://github.com/abjadai/catt/releases/download/v2/best_ed_mlm_ns_epoch_178.pt
wget -P models/ https://github.com/abjadai/catt/releases/download/v2/best_eo_mlm_ns_epoch_193.pt

You can use the inference code examples: predict_ed.py for ED models and predict_eo.py for EO models.
Both examples are provided with batch inference support. Read the source code to gain a better understanding.

python predict_ed.py
python predict_eo.py

EO models are recommended for faster inference.
ED models are recommended for better accuracy of the predicted diacritics.

To convert your trained PyTorch models to ONNX format, use the export script:

python export_to_onnx.py

This script will:

• Load your trained PyTorch model checkpoints
• Export separate ONNX models for encoder and decoder components
• Validate the exported models for correctness
• Save the ONNX models in the onnx_models/ directory

Output files:

• encoder.onnx - The encoder component
• decoder.onnx - The decoder component (or linear layer for encoder-only models)

To test and run inference with the exported ONNX models:

python run_onnx.py

This script will:

• Load the exported ONNX models
• Run inference using ONNX Runtime

For more details on the export process, check the export_to_onnx.py script configuration.

To start trainnig, you need to download the dataset from the Releases section of this repo.

wget https://github.com/abjadai/catt/releases/download/v2/dataset.zip
unzip dataset.zip

Then, edit the script train_catt.py and adjest the default values:

# Model's Configs
model_type = 'ed' # 'eo' for Encoder-Only OR 'ed' for Encoder-Decoder
dl_num_workers = 32
batch_size = 32
max_seq_len = 1024
threshold = 0.6

# Pretrained Char-Based BERT
pretrained_mlm_pt = None # Use None if you want to initialize weights randomly OR the path to the char-based BERT
#pretrained_mlm_pt = 'char_bert_model_pretrained.pt'

Finally, run the training script.

python train_catt.py

• This code is mainly adapted from this repo.
• An older version of some Arabic scripts that are available in pyarabic library were used as well.

• inference script
• upload our pretrained models
• upload CATT dataset
• upload DER scripts
• training script

This repository has updated its license from Creative Commons Attribution-NonCommercial (CC-BY-NC) to Apache 2.0 License. This change removes commercial use restrictions and adopts an industry-standard open-source license, enabling broader adoption and collaboration. This transition supports the project's growth while maintaining our commitment to open-source principles.