CATT & CATT-Whisper (code, dataset, models, papers, spaces)

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Abjad[[:space:]]AI[[:space:]]at[[:space:]]NADI[[:space:]]2025.[[:space:]]CATT-Whisper.[[:space:]]Multimodal[[:space:]]Diacritic[[:space:]]Restoration[[:space:]]Using[[:space:]]Text[[:space:]]and[[:space:]]Speech[[:space:]]Representations.pdf filter=lfs diff=lfs merge=lfs -text
+CATT.[[:space:]]Character-based[[:space:]]Arabic[[:space:]]Tashkeel[[:space:]]Transformer.pdf filter=lfs diff=lfs merge=lfs -text
+NADI[[:space:]]2025.[[:space:]]The[[:space:]]First[[:space:]]Multidialectal[[:space:]]Arabic[[:space:]]Speech[[:space:]]Processing[[:space:]]Shared[[:space:]]Task.pdf filter=lfs diff=lfs merge=lfs -text
+Thaka[[:space:]]at[[:space:]]KSAA-2026[[:space:]]Task[[:space:]]2.[[:space:]]Regularized[[:space:]]Fine-Tuning[[:space:]]for[[:space:]]Arabic[[:space:]]Speech[[:space:]]Diacritization.pdf filter=lfs diff=lfs merge=lfs -text

Abjad AI at NADI 2025. CATT-Whisper. Multimodal Diacritic Restoration Using Text and Speech Representations.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16235cec951d5a337c988d5fa66b809946d070ea69244f16242d93f49057cc48
+size 296424

CATT. Character-based Arabic Tashkeel Transformer.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbbef8cd0bfa0f7a721af845316e82da4bcc735635b079127933d14455e14610
+size 400230

NADI 2025. The First Multidialectal Arabic Speech Processing Shared Task.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:633ffcd7c2dcdb1b3c9e7c16e0ec9c0850285f313c882e076e286e684b1ba45a
+size 1215064

Thaka at KSAA-2026 Task 2. Regularized Fine-Tuning for Arabic Speech Diacritization.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b063a23f75d4104f1a34087ab00b81a4641fe2195edf3936e5624dc854917ad5
+size 159709

code/Arabic-Romanizer (ArabicTransliterator +8).zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98fed58e6d6f1fad9f245a5f915061bcd11f69a85f193da2dde63aa3ad83e853
+size 22718916

code/ArabicTransliterator.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81ef1cbdc5c266d6e501010b39f9998ef3d3a03dec095a5c722fe2d0c9b679ec
+size 46528355

code/catt-whisper.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:220f178a205ca4bfff5f4cbf752979dc273e38a6b29f724c045d84945d5a9e9e
+size 1271147

code/catt.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93f4f236ba71de710cf2316558f8bf0b0d3e5853c84a13659f8505c13bc90771
+size 8552916

code/mishkal.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9243591c7a437ee16dccc343d51e61e3d1b6b2bd9383a42863d95c9ec108495
+size 117913533

code/saara.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8fcb5eef9e6ad428cb1129481175b1af32009a072e8f353d2eb79d10460a059
+size 46530545

dataset/dataset.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa332ca9d3118ed8bc974debc681ba3f9d5ebca3cb824073b24bef0314e5fb00
+size 209268505

models/CATT-Whisper/catt_whisper_base_model_v1_epoch_26_with_spec_augment.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7ca2e8107eaab302b2c48a91f4d4f3ec8c151f96ee5ffd27750778a51f5ecac
+size 159223099

models/CATT-Whisper/source.txt

@@ -0,0 +1 @@
+https://github.com/abjadai/catt-whisper

models/CATT/best_ed_epoch_175.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5125495c8fd6c5480c8443a2d73b766be00c1c85988d6b67ec4e5dc7c17bdd5a
+size 88403216

models/CATT/best_ed_mlm_epoch_175.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a250c29acdb8cbdd9e729287a6631710986d8243d487b2fa5a25b996cbaa84be
+size 88403608

models/CATT/best_ed_mlm_ns_epoch_178.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6297ec54f9af518a59d4b9e5831d3ad4da304435b8c7fb969f6e3a16816ac23
+size 88403510

models/CATT/best_eo_epoch_192.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dbe3b98ca741382fbdb7acdf5a2314b7a93bc295b2e3364c6343b9a7c8985081
+size 75761995

models/CATT/best_eo_mlm_epoch_192.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f21dfe38acf867b1531be3e20ab76c6b2e438eaae6e1ff108607be986518d11
+size 75762311

models/CATT/best_eo_mlm_ns_epoch_193.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ad3e811b1bf7ecda252dc9284d2bf3087c1795add47be5ff34cb13688413322
+size 75762232

models/CATT/char_bert_model_pretrained.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:440b0dcad32278a607d34d110a00a0ff20d9ad7d67f7dea8b5e323645976c9c4
+size 75807234

models/CATT/ed_epoch_5.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47b811b4db7675876a83356af04c1f61b2d60aab4ce57ebfbe848aa95f4b7a05
+size 88402530

models/CATT/ed_mlm_epoch_5.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f5534cf55306e31ad7cbdd89f931a1ae1b4724f04d6d68391b4e699f52fc1df
+size 88402922

models/CATT/ed_model_onnx.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75d2a6360725c7e0bbb238b5a950f621fb8e30b21ad7c65e91e0f36759bcb3ce
+size 85926550

models/CATT/eo_epoch_5.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0997479729f9cbbfa0b10070e0099fcc88ff540e9dcd113b2e0c27bd59f0ca8b
+size 75761442

models/CATT/eo_mlm_epoch_5.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2bd85659b985949849c7377ddd941f3965bec60d0b8eb60f5a012f10f46bec3
+size 75761758

models/CATT/eo_model_onnx.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5343bfdfda70c981b587811b79348b8dde91f65c0a551ca98f232ea5cb62ceac
+size 72109998

models/CATT/source.txt

@@ -0,0 +1 @@
+https://github.com/abjadai/catt

models/catt-diacritizer/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

models/catt-diacritizer/README.md

@@ -0,0 +1,43 @@
+---
+license: apache-2.0
+language: ar
+library_name: onnx
+tags:
+- arabic
+- diacritization
+- tashkeel
+- onnx
+- transformer
+---
+
+# CATT-EO — Arabic diacritizer (stitched single ONNX)
+
+Single-file ONNX of the **encoder-only [CATT](https://github.com/abjadai/catt)** model
+(Abjad AI, Apache-2.0). CATT-EO is non-autoregressive — a transformer encoder followed by
+a linear classifier head — so [TigreGotico](https://huggingface.co/TigreGotico)
+**stitched the two stages into one ONNX** with `onnx.compose`. The stitched graph is
+**byte-identical** (max|Δ| = 0.0) to the original two-stage encoder→head pipeline.
+
+| file | precision | size |
+|------|-----------|-----:|
+| `catt_eo.onnx` | fp32 | 78 MB |
+| `catt_eo.int8.onnx` | dynamic int8 | 21.6 MB |
+
+**I/O.** Inputs `src` (int64, `[batch, seq]` Buckwalter ids) and `src_mask`
+(bool, `[batch, 1, seq, seq]`). Output `[batch, seq, 18]` tag logits. Tokenizer:
+`tashkeel_tokenizer_onnx.py` (+ `bw2ar.py`, `utils.py`) — Buckwalter encoding and the
+18-tag tashkeel scheme. See `catt_models_onnx.py` for an end-to-end example.
+
+**Benchmark.** ~4.27 % DER on the broad
+[`TigreGotico/arabic_diacritized_text`](https://huggingface.co/datasets/TigreGotico/arabic_diacritized_text)
+test (much lower on CATT's own narrow benchmark — distribution-dependent). The int8 file
+is dynamic-quantized; expect small divergence from fp32.
+
+Part of the
+[Arabic Diacritizers / Tashkeel collection](https://huggingface.co/collections/TigreGotico/arabic-diacritizers-tashkeel-6a247318559bcc49e128aa5f).
+Original CATT © Abjad AI (Apache-2.0); this is a stitched re-export with attribution.
+
+
+## Encoder-decoder variant (CATT-ED)
+
+`catt_ed_encoder.onnx` + `catt_ed_decoder.onnx` (+ `.int8`) are the **autoregressive** encoder-decoder model — the decoder runs once per output position (causal mask), so it is two ONNX (not stitchable) and much slower than CATT-EO. In `text2tashkeel` it is the `catt-ed` / `catt-ed-int8` model.

models/catt-diacritizer/bw2ar.py

@@ -0,0 +1,126 @@
+"""
+functions to convert Arabic words/text into buckwalter encoding and vice versa
+"""
+
+import sys
+import re
+import utils
+
+buck2uni = {
+            "'": u"\u0621", # hamza-on-the-line
+            "|": u"\u0622", # madda
+            ">": u"\u0623", # hamza-on-'alif
+            "&": u"\u0624", # hamza-on-waaw
+            "<": u"\u0625", # hamza-under-'alif
+            "}": u"\u0626", # hamza-on-yaa'
+            "A": u"\u0627", # bare 'alif
+            "b": u"\u0628", # baa'
+            "p": u"\u0629", # taa' marbuuTa
+            "t": u"\u062A", # taa'
+            "v": u"\u062B", # thaa'
+            "j": u"\u062C", # jiim
+            "H": u"\u062D", # Haa'
+            "x": u"\u062E", # khaa'
+            "d": u"\u062F", # daal
+            "*": u"\u0630", # dhaal
+            "r": u"\u0631", # raa'
+            "z": u"\u0632", # zaay
+            "s": u"\u0633", # siin
+            "$": u"\u0634", # shiin
+            "S": u"\u0635", # Saad
+            "D": u"\u0636", # Daad
+            "T": u"\u0637", # Taa'
+            "Z": u"\u0638", # Zaa' (DHaa')
+            "E": u"\u0639", # cayn
+            "g": u"\u063A", # ghayn
+            "_": u"\u0640", # taTwiil
+            "f": u"\u0641", # faa'
+            "q": u"\u0642", # qaaf
+            "k": u"\u0643", # kaaf
+            "l": u"\u0644", # laam
+            "m": u"\u0645", # miim
+            "n": u"\u0646", # nuun
+            "h": u"\u0647", # haa'
+            "w": u"\u0648", # waaw
+            "Y": u"\u0649", # 'alif maqSuura
+            "y": u"\u064A", # yaa'
+            "F": u"\u064B", # fatHatayn
+            "N": u"\u064C", # Dammatayn
+            "K": u"\u064D", # kasratayn
+            "a": u"\u064E", # fatHa
+            "u": u"\u064F", # Damma
+            "i": u"\u0650", # kasra
+            "~": u"\u0651", # shaddah
+            "o": u"\u0652", # sukuun
+            "`": u"\u0670", # dagger 'alif
+            "{": u"\u0671", # waSla
+}
+
+# For a reverse transliteration (Unicode -> Buckwalter), a dictionary
+# which is the reverse of the above buck2uni is essential.
+uni2buck = {}
+
+# Iterate through all the items in the buck2uni dict.
+for (key, value) in buck2uni.items():
+    # The value from buck2uni becomes a key in uni2buck, and vice
+    # versa for the keys.
+    uni2buck[value] = key
+
+# add special characters
+uni2buck[u"\ufefb"] = "lA"
+uni2buck[u"\ufef7"] = "l>"
+uni2buck[u"\ufef5"] = "l|"
+uni2buck[u"\ufef9"] = "l<"
+
+# clean the arabic text from unwanted characters that may cause problem while building the language model
+def clean_text(text):
+    text = re.sub(u"[\ufeff]", "", text,  flags=re.UNICODE) # strip Unicode Character 'ZERO WIDTH NO-BREAK SPACE' (U+FEFF). For more info, check http://www.fileformat.info/info/unicode/char/feff/index.htm
+    text = utils.remove_non_arabic(text)
+    text = utils.strip_tashkeel(text)
+    text = utils.strip_tatweel(text)
+    return text
+
+# convert a single word into buckwalter and vice versa
+def transliterate_word(input_word, direction='bw2ar'):
+    output_word = ''
+    # Loop over each character in the string, bw_word.
+    for char in input_word:
+        # Look up current char in the dictionary to get its
+        # respective value. If there is no match, e.g., chars like
+        # spaces, then just stick with the current char without any
+        # conversion.
+        # if type(char) == bytes:
+        #    char = char.decode('ascii')
+        if direction == 'bw2ar':
+            #print('in bw2ar')
+            output_word += buck2uni.get(char, char)
+        elif direction == 'ar2bw':
+            #print('in ar2bw')
+            output_word += uni2buck.get(char, char)
+        else:
+            sys.stderr.write('Error: invalid direction!')
+            sys.exit()
+    return output_word
+
+
+# convert a text into buckwalter and vice versa
+def transliterate_text(input_text, direction='bw2ar'):
+    output_text = ''
+    for input_word in input_text.split(' '):
+        output_text += transliterate_word(input_word, direction) + ' '
+
+    return output_text[:-1] # remove the last space ONLY
+
+
+if __name__ == '__main__':
+    if len(sys.argv) < 2:
+        sys.stderr.write('Usage: INPUT TEXT | python {} DIRECTION(bw2ar|ar2bw)'.format(sys.argv[1]))
+        exit(1)
+    for line in sys.stdin:
+        line = line if sys.argv[1] == 'bw2ar' else clean_text(line)
+        output_text = transliterate_text(line, direction=str(sys.argv[1]))
+        if output_text.strip() != '':
+            sys.stdout.write('{}\n'.format(output_text.strip()))
+
+
+

models/catt-diacritizer/catt_ed_decoder.int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5819a139e8b27ed2bd8d0209b6a75206e86f1a3d3dc1b359f75b4fcbcfa2e4e
+size 15158542

models/catt-diacritizer/catt_ed_decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14b5580ecb33f8687cdccb14edb120afa2a2cfbf7caf4444fa4a2a86548c375c
+size 52991044

models/catt-diacritizer/catt_ed_encoder.int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d5dff46db736a1b1c17babad46db499322cd0c30da65121fd45446b03e806db8
+size 11845630

models/catt-diacritizer/catt_ed_encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ec637243feb55f3b83b35264a5daa9c23e1261097ce42ecc85163fba5e1d962
+size 40235317

models/catt-diacritizer/catt_eo.int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59ba20d02a4b507a6b816a9b46b1230ba73d341c6731aa48da5de99388a39fca
+size 21564456

models/catt-diacritizer/catt_eo.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f4040ae643c2a5aa40221bb845e146d8826690a95501c9250f72dbc5a12479c
+size 78308696

models/catt-diacritizer/catt_models_onnx.py

@@ -0,0 +1,273 @@
+import os
+import math
+import numpy as np
+from tqdm import tqdm
+from pathlib import Path
+import onnxruntime as ort
+from abc import ABC, abstractmethod
+from utils import remove_non_arabic
+from tashkeel_tokenizer_onnx import TashkeelTokenizer
+
+
+class BaseONNXTashkeel(ABC):
+    """
+    Base class for ONNX-based Arabic text diacritization models
+
+    This abstract base class provides common functionality for Arabic text
+    diacritization using pre-trained ONNX models.
+    """
+
+    def __init__(self, encoder_path=None, decoder_path=None, tokenizer=None,
+                 auto_preprocess=True, model_type="ed_model"):
+        """
+        Initialize ONNX Tashkeel model
+
+        Args:
+            encoder_path (str, optional): Path to encoder ONNX file
+            decoder_path (str, optional): Path to decoder ONNX file
+            tokenizer (TashkeelTokenizer, optional): Tokenizer instance
+            auto_preprocess (bool): Automatically preprocess text (remove non-Arabic)
+            model_type (str): Type of model ("ed_model" or "eo_model")
+        """
+        self.tokenizer = tokenizer or TashkeelTokenizer()
+        self.auto_preprocess = auto_preprocess
+        self.src_pad_idx = self.tokenizer.letters_map['<PAD>']
+
+        # Set model-specific paths
+        encoder_path, decoder_path = self._get_model_paths(
+            encoder_path, decoder_path, model_type
+        )
+
+        # Load ONNX sessions
+        self._load_onnx_sessions(encoder_path, decoder_path)
+
+    def _get_model_paths(self, encoder_path, decoder_path, model_type):
+        """Get default model paths if not provided"""
+        if encoder_path is None or decoder_path is None:
+            package_dir = Path(__file__).parent
+            model_dir = package_dir / "onnx_models" / model_type
+            encoder_path = encoder_path or str(model_dir / "encoder.onnx")
+            decoder_path = decoder_path or str(model_dir / "decoder.onnx")
+        return encoder_path, decoder_path
+
+    def _load_onnx_sessions(self, encoder_path, decoder_path):
+        """Load ONNX runtime sessions"""
+        if not os.path.exists(encoder_path):
+            raise FileNotFoundError(f"Encoder ONNX file not found: {encoder_path}")
+        if not os.path.exists(decoder_path):
+            raise FileNotFoundError(f"Decoder ONNX file not found: {decoder_path}")
+
+        _opt = ort.SessionOptions()
+        import os as _os
+        _t = _os.environ.get("TT_ORT_THREADS")
+        if _t:
+            _opt.intra_op_num_threads = int(_t); _opt.inter_op_num_threads = int(_t)
+        self.encoder_session = ort.InferenceSession(encoder_path, sess_options=_opt)
+        self.decoder_session = ort.InferenceSession(decoder_path, sess_options=_opt)
+        print("✅ Loaded model successfully!")
+
+    def _preprocess_text(self, text):
+        """Preprocess text if auto_preprocess is enabled"""
+        return remove_non_arabic(text) if self.auto_preprocess else text
+
+    def _make_pad_mask(self, q, k, q_pad_idx, k_pad_idx):
+        """Create padding mask (numpy version)"""
+        len_q, len_k = q.shape[1], k.shape[1]
+        k_mask = (k != k_pad_idx).reshape(-1, 1, 1, len_k)
+        k_mask = np.repeat(k_mask, len_q, axis=2)
+        q_mask = (q != q_pad_idx).reshape(-1, 1, len_q, 1)
+        q_mask = np.repeat(q_mask, len_k, axis=3)
+        return k_mask & q_mask
+
+    def _make_no_peak_mask(self, q, k):
+        """Create no-peak mask (numpy version)"""
+        len_q, len_k = q.shape[1], k.shape[1]
+        return np.tril(np.ones((len_q, len_k))).astype(bool)
+
+    def _run_encoder(self, src_np):
+        """Run encoder using ONNX"""
+        src_mask_np = self._make_pad_mask(
+            src_np, src_np, self.src_pad_idx, self.src_pad_idx
+        ).astype(bool)
+
+        encoder_inputs = {'src': src_np, 'src_mask': src_mask_np}
+        return self.encoder_session.run(None, encoder_inputs)[0]
+
+    def _pad_seq(self, sequences, batch_first=True, padding_value=0.0, prepadding=True):
+        """Pad sequences to same length"""
+        lens = [seq.shape[0] for seq in sequences]
+        max_len = max(lens)
+
+        # Create padded array
+        if len(sequences[0].shape) == 1:
+            padded_sequences = np.full((len(sequences), max_len), padding_value, dtype=sequences[0].dtype)
+        else:
+            shape = (len(sequences), max_len) + sequences[0].shape[1:]
+            padded_sequences = np.full(shape, padding_value, dtype=sequences[0].dtype)
+
+        # Fill with actual sequence data
+        for i, seq in enumerate(sequences):
+            padded_sequences[i, :lens[i]] = seq
+
+        if prepadding:
+            for i in range(len(lens)):
+                padded_sequences[i] = np.roll(padded_sequences[i], max_len - lens[i], axis=0)
+
+        if not batch_first:
+            padded_sequences = padded_sequences.transpose(1, 0, *range(2, padded_sequences.ndim))
+
+        return padded_sequences
+
+    def _get_batches(self, X, batch_size=16):
+        """Generate batches from list"""
+        num_batches = math.ceil(len(X) / batch_size)
+        for i in range(num_batches):
+            yield X[i*batch_size : (i+1)*batch_size]
+
+    def _prepare_batch_input(self, texts_mini):
+        """Prepare input batch for processing"""
+        input_ids_list = []
+        for text in texts_mini:
+            input_ids, _ = self.tokenizer.encode(text, test_match=False)
+            input_ids_list.append(input_ids)
+
+        return self._pad_seq(
+            input_ids_list,
+            batch_first=True,
+            padding_value=self.tokenizer.letters_map['<PAD>'],
+            prepadding=False
+        )
+
+    def _apply_space_mask(self, predictions, input_ids):
+        """Apply space masking to predictions"""
+        space_mask = (input_ids == self.tokenizer.letters_map[' '])
+        predictions[space_mask] = self.tokenizer.tashkeel_map[self.tokenizer.no_tashkeel_tag]
+        return predictions
+
+    @abstractmethod
+    def _process_batch(self, batch_input_ids):
+        """Process a batch of input IDs - to be implemented by subclasses"""
+        pass
+
+    def do_tashkeel_batch(self, texts, batch_size=16, verbose=True):
+        """
+        Add tashkeel to multiple Arabic texts
+
+        Args:
+            texts (list): List of input Arabic texts
+            batch_size (int): Batch size for processing
+            verbose (bool): Show progress bar
+
+        Returns:
+            list: List of texts with tashkeel
+        """
+        processed_texts = [self._preprocess_text(text) for text in texts]
+        text_with_tashkeel = []
+        data_iter = self._get_batches(processed_texts, batch_size)
+
+        if verbose:
+            num_batches = math.ceil(len(processed_texts) / batch_size)
+            data_iter = tqdm(data_iter, total=num_batches)
+
+        for texts_mini in data_iter:
+            batch_input_ids = self._prepare_batch_input(texts_mini)
+            predictions = self._process_batch(batch_input_ids)
+            text_with_tashkeel_mini = self.tokenizer.decode(batch_input_ids, predictions)
+            text_with_tashkeel.extend(text_with_tashkeel_mini)
+
+        return text_with_tashkeel
+
+    def do_tashkeel(self, text, verbose=True):
+        """
+        Add tashkeel to Arabic text
+
+        Args:
+            text (str): Input Arabic text
+
+        Returns:
+            str: Text with tashkeel
+        """
+        return self.do_tashkeel_batch([text], verbose=False)[0]
+
+
+class CATTEncoderDecoder(BaseONNXTashkeel):
+    """
+    ONNX-based Arabic text diacritization model using encoder-decoder architecture
+    """
+
+    def __init__(self, encoder_path=None, decoder_path=None, tokenizer=None,
+                 auto_preprocess=True):
+        """Initialize encoder-decoder model"""
+        super().__init__(encoder_path, decoder_path, tokenizer, auto_preprocess, "ed_model")
+        self.trg_pad_idx = self.tokenizer.tashkeel_map['<PAD>']
+
+    def _run_decoder(self, trg_np, enc_src_np, src_np):
+        """Run decoder using ONNX"""
+        src_trg_mask_np = self._make_pad_mask(
+            trg_np, src_np, self.trg_pad_idx, self.src_pad_idx
+        ).astype(bool)
+
+        trg_mask_np = self._make_pad_mask(
+            trg_np, trg_np, self.trg_pad_idx, self.trg_pad_idx
+        ).astype(bool)
+
+        no_peak_mask_np = self._make_no_peak_mask(trg_np, trg_np)
+        trg_mask_np = trg_mask_np & no_peak_mask_np
+
+        decoder_inputs = {
+            'trg': trg_np,
+            'enc_src': enc_src_np,
+            'trg_mask': trg_mask_np,
+            'src_trg_mask': src_trg_mask_np
+        }
+
+        return self.decoder_session.run(None, decoder_inputs)[0]
+
+    def _process_batch(self, batch_input_ids):
+        """Process batch using encoder-decoder with autoregressive decoding"""
+        src = batch_input_ids
+        enc_src = self._run_encoder(src)
+
+        # Initialize target with BOS token
+        target_ids = np.array([[self.tokenizer.tashkeel_map['<BOS>']]] * len(batch_input_ids), dtype=np.int64)
+
+        # Autoregressive decoding loop
+        for i in range(src.shape[1] - 1):
+            preds = self._run_decoder(target_ids, enc_src, src)
+            next_tokens = np.argmax(preds[:, -1, :], axis=1)
+            target_ids = np.concatenate([target_ids, np.expand_dims(next_tokens, 1)], axis=1)
+
+            # Apply space handling logic
+            target_ids = self._apply_space_mask(target_ids, src[:, :target_ids.shape[1]])
+
+        return target_ids
+
+
+class CATTEncoderOnly(BaseONNXTashkeel):
+    """
+    ONNX-based Arabic text diacritization model using encoder-only architecture
+    """
+
+    def __init__(self, encoder_path=None, decoder_path=None, tokenizer=None,
+                 auto_preprocess=True):
+        """Initialize encoder-only model"""
+        super().__init__(encoder_path, decoder_path, tokenizer, auto_preprocess, "eo_model")
+
+    def _run_decoder(self, enc_src_np):
+        """Run decoder using ONNX"""
+        decoder_inputs = {'enc_src': enc_src_np}
+        return self.decoder_session.run(None, decoder_inputs)[0]
+
+    def _process_batch(self, batch_input_ids):
+        """Process batch using encoder-only architecture"""
+        # Remove BOS and EOS tokens
+        batch_input_ids = batch_input_ids[:, 1:-1]
+
+        enc_src = self._run_encoder(batch_input_ids)
+        preds = self._run_decoder(enc_src)
+        predictions = np.argmax(preds, axis=-1)
+
+        # Apply space masking
+        predictions = self._apply_space_mask(predictions, batch_input_ids)
+
+        return predictions

models/catt-diacritizer/source.txt

@@ -0,0 +1 @@
+https://huggingface.co/TigreGotico/catt-diacritizer

models/catt-diacritizer/tashkeel_tokenizer_onnx.py

@@ -0,0 +1,158 @@
+
+import re
+import numpy as np
+import bw2ar
+from utils import HARAKAT_PAT
+
+
+class TashkeelTokenizer:
+
+    def __init__(self):
+        self.letters = [' ', '$', '&', "'", '*', '<', '>', 'A', 'D', 'E', 'H', 'S', 'T', 'Y', 'Z',
+                        'b', 'd', 'f', 'g', 'h', 'j', 'k', 'l', 'm', 'n', 'p', 'q', 'r', 's', 't',
+                        'v', 'w', 'x', 'y', 'z', '|', '}'
+                       ]
+        self.letters = ['<PAD>', '<BOS>', '<EOS>'] + self.letters + ['<MASK>']
+
+        self.no_tashkeel_tag = '<NT>'
+        self.tashkeel_list = ['<NT>', '<SD>', '<SDD>', '<SF>', '<SFF>', '<SK>',
+                               '<SKK>', 'F', 'K', 'N', 'a', 'i', 'o', 'u', '~']
+
+        self.tashkeel_list = ['<PAD>', '<BOS>', '<EOS>'] + self.tashkeel_list
+
+        self.tashkeel_map = {c:i for i,c in enumerate(self.tashkeel_list)}
+        self.letters_map = {c:i for i,c in enumerate(self.letters)}
+        self.inverse_tags = {
+                 '~a': '<SF>',  # shaddah and fatHa
+                 '~u': '<SD>',  # shaddah and Damma
+                 '~i': '<SK>',  # shaddah and kasra
+                 '~F': '<SFF>', # shaddah and fatHatayn
+                 '~N': '<SDD>', # shaddah and Dammatayn
+                 '~K': '<SKK>'  # shaddah and kasratayn
+        }
+        self.tags = {v:k for k,v in self.inverse_tags.items()}
+        self.shaddah_last  = ['a~', 'u~', 'i~', 'F~', 'N~', 'K~']
+        self.shaddah_first = ['~a', '~u', '~i', '~F', '~N', '~K']
+        self.tahkeel_chars = ['F','N','K','a', 'u', 'i', '~', 'o']
+
+    def clean_text(self, text):
+        text = re.sub(u'[%s]' % u'\u0640', '', text) # strip tatweel
+        text = text.replace('ٱ', 'ا')
+        return ' '.join(re.sub(u"[^\u0621-\u063A\u0640-\u0652\u0670\u0671\ufefb\ufef7\ufef5\ufef9 ]", " ", text,  flags=re.UNICODE).split())
+
+    def check_match(self, text_with_tashkeel, letter_n_tashkeel_pairs):
+        text_with_tashkeel = text_with_tashkeel.strip()
+        # test if the reconstructed text with tashkeel is the same as the original one
+        syn_text = self.combine_tashkeel_with_text(letter_n_tashkeel_pairs)
+        return syn_text == text_with_tashkeel or syn_text == self.unify_shaddah_position(text_with_tashkeel)
+
+    def unify_shaddah_position(self, text_with_tashkeel):
+        # unify the order of shaddah and the harakah to make shaddah always at the beginning
+        for i in range(len(self.shaddah_first)):
+            text_with_tashkeel = text_with_tashkeel.replace(self.shaddah_last[i], self.shaddah_first[i])
+        return text_with_tashkeel
+
+    def split_tashkeel_from_text(self, text_with_tashkeel, test_match=True):
+        text_with_tashkeel = self.clean_text(text_with_tashkeel)
+        text_with_tashkeel = bw2ar.transliterate_text(text_with_tashkeel, 'ar2bw')
+        text_with_tashkeel = text_with_tashkeel.replace('`', '') # remove dagger 'alif
+
+        # unify the order of shaddah and the harakah to make shaddah always at the beginning
+        text_with_tashkeel = self.unify_shaddah_position(text_with_tashkeel)
+
+        # remove duplicated harakat
+        for i in range(len(self.tahkeel_chars)):
+            text_with_tashkeel = text_with_tashkeel.replace(self.tahkeel_chars[i]*2, self.tahkeel_chars[i])
+
+        letter_n_tashkeel_pairs = []
+        for i in range(len(text_with_tashkeel)): # go over the whole text
+            # check if the first character is a normal letter and the second character is a tashkeel
+            if i < (len(text_with_tashkeel) - 1) and not text_with_tashkeel[i] in self.tashkeel_list and text_with_tashkeel[i+1] in self.tashkeel_list:
+                # IMPORTANT: check if tashkeel is Shaddah, then there might be another Tashkeel char associated with it. If so,
+                # replace both Shaddah and the Tashkeel chars with the appropriate tag
+                if text_with_tashkeel[i+1] == '~':
+                    # IMPORTANT: the following if statement depends on the concept of short circuit!!
+                    # The first condition checks if there are still more chars before it access position i+2
+                    # "text_with_tashkeel[i+2]" since it causes "index out of range" exception. Notice that
+                    # Shaddah here is put in the first position before the Harakah.
+                    if i+2 < len(text_with_tashkeel) and f'~{text_with_tashkeel[i+2]}' in self.inverse_tags:
+                        letter_n_tashkeel_pairs.append((text_with_tashkeel[i], self.inverse_tags[f'~{text_with_tashkeel[i+2]}']))
+                    else:
+                        # if it is only Shaddah, just add it to the list
+                        letter_n_tashkeel_pairs.append((text_with_tashkeel[i], '~'))
+                else:
+                    letter_n_tashkeel_pairs.append((text_with_tashkeel[i], text_with_tashkeel[i+1]))
+            # if the character at position i is a normal letter and has no Tashkeel, then add
+            # it with the tag "self.no_tashkeel_tag"
+            # IMPORTANT: this elif block ensures also that there is no two or more consecutive tashkeel other than shaddah
+            elif not text_with_tashkeel[i] in self.tashkeel_list:
+                letter_n_tashkeel_pairs.append((text_with_tashkeel[i], self.no_tashkeel_tag))
+
+        if test_match:
+            # test if the split is done correctly by ensuring that we can retrieve back the original text
+            assert self.check_match(text_with_tashkeel, letter_n_tashkeel_pairs)
+        return [('<BOS>', '<BOS>')] + letter_n_tashkeel_pairs + [('<EOS>', '<EOS>')]
+
+    def combine_tashkeel_with_text(self, letter_n_tashkeel_pairs):
+        combined_with_tashkeel = []
+        for letter, tashkeel in letter_n_tashkeel_pairs:
+            combined_with_tashkeel.append(letter)
+            if tashkeel in self.tags:
+                combined_with_tashkeel.append(self.tags[tashkeel])
+            elif tashkeel != self.no_tashkeel_tag:
+                combined_with_tashkeel.append(tashkeel)
+        text = ''.join(combined_with_tashkeel)
+        return text
+
+    def encode(self, text_with_tashkeel, test_match=True):
+        letter_n_tashkeel_pairs = self.split_tashkeel_from_text(text_with_tashkeel, test_match)
+        text, tashkeel = zip(*letter_n_tashkeel_pairs)
+        input_ids = [self.letters_map[c] for c in text]
+        target_ids = [self.tashkeel_map[c] for c in tashkeel]
+        return np.array(input_ids, dtype=np.int64), np.array(target_ids, dtype=np.int64)
+
+    def filter_tashkeel(self, tashkeel):
+        tmp = []
+        for i, t in enumerate(tashkeel):
+            if i != 0 and t == '<BOS>':
+                t = self.no_tashkeel_tag
+            elif i != (len(tashkeel) - 1) and t == '<EOS>':
+                t = self.no_tashkeel_tag
+            tmp.append(t)
+        tashkeel = tmp
+        return tashkeel
+
+    def decode(self, input_ids, target_ids):
+        ar_texts = []
+        for j in range(len(input_ids)):
+            letters = [self.letters[i] for i in input_ids[j]]
+            tashkeel = [self.tashkeel_list[i] for i in target_ids[j]]
+
+            letters = list(filter(lambda x: x != '<BOS>' and x != '<EOS>' and x != '<PAD>', letters))
+            tashkeel = self.filter_tashkeel(tashkeel)
+            tashkeel = list(filter(lambda x: x != '<BOS>' and x != '<EOS>' and x != '<PAD>', tashkeel))
+
+            # VERY IMPORTANT NOTE: zip takes min(len(letters), len(tashkeel)) and discard the reset of letters / tashkeels
+            letter_n_tashkeel_pairs = list(zip(letters, tashkeel))
+            bw_text = self.combine_tashkeel_with_text(letter_n_tashkeel_pairs)
+            ar_text = bw2ar.transliterate_text(bw_text, 'bw2ar')
+            ar_texts.append(ar_text)
+        return ar_texts
+
+    def get_tashkeel_with_case_ending(self, text, case_ending=True):
+        text_split = self.split_tashkeel_from_text(text, test_match=False)
+        text_spaces_indecies = [i for i, el in enumerate(text_split) if el == (' ', '<NT>')]
+        new_text_split = []
+        for i, el in enumerate(text_split):
+            if not case_ending and (i+1) in text_spaces_indecies:
+                el = (el[0], '<NT>') # no case ending
+            new_text_split.append(el)
+        letters, tashkeel = zip(*new_text_split)
+        return letters, tashkeel
+
+    def remove_tashkeel(self, text):
+        text = HARAKAT_PAT.sub('', text)
+        text = re.sub(u"[\u064E]", "", text,  flags=re.UNICODE) # fattha
+        text = re.sub(u"[\u0671]", "", text,  flags=re.UNICODE) # waSla
+        return text
+

models/catt-diacritizer/utils.py

@@ -0,0 +1,139 @@
+"""
+note: this code is used in bw2ar.py file
+"""
+
+#!/usr/bin/python
+# -*- coding=utf-8 -*-
+#---
+# $Id: arabic.py,v 1.6 2003/04/22 17:18:22 elzubeir Exp $
+#
+# ------------
+# Description:
+# ------------
+#
+# Arabic codes
+#
+# (C) Copyright 2003, Arabeyes, Mohammed Elzubeir
+# (C) Copyright 2019, Faris Abdullah Alasmary
+# -----------------
+# Revision Details:    (Updated by Revision Control System)
+# -----------------
+#  $Date: 2003/04/22 17:18:22 $
+#  $Author: elzubeir $
+#  $Revision: 1.6 $
+#  $Source: /home/arabeyes/cvs/projects/duali/pyduali/pyduali/arabic.py,v $
+#
+#  This program is written under the BSD License.
+#---
+""" Constants for arabic """
+import re
+COMMA = u'\u060C'
+SEMICOLON = u'\u061B'
+QUESTION = u'\u061F'
+HAMZA = u'\u0621'
+ALEF_MADDA = u'\u0622'
+ALEF_HAMZA_ABOVE = u'\u0623'
+WAW_HAMZA = u'\u0624'
+ALEF_HAMZA_BELOW = u'\u0625'
+YEH_HAMZA = u'\u0626'
+ALEF = u'\u0627'
+BEH = u'\u0628'
+TEH_MARBUTA = u'\u0629'
+TEH = u'\u062a'
+THEH = u'\u062b'
+JEEM = u'\u062c'
+HAH = u'\u062d'
+KHAH = u'\u062e'
+DAL = u'\u062f'
+THAL = u'\u0630'
+REH = u'\u0631'
+ZAIN = u'\u0632'
+SEEN = u'\u0633'
+SHEEN = u'\u0634'
+SAD = u'\u0635'
+DAD = u'\u0636'
+TAH = u'\u0637'
+ZAH = u'\u0638'
+AIN = u'\u0639'
+GHAIN = u'\u063a'
+TATWEEL = u'\u0640'
+FEH = u'\u0641'
+QAF = u'\u0642'
+KAF = u'\u0643'
+LAM = u'\u0644'
+MEEM = u'\u0645'
+NOON = u'\u0646'
+HEH = u'\u0647'
+WAW = u'\u0648'
+ALEF_MAKSURA = u'\u0649'
+YEH = u'\u064a'
+MADDA_ABOVE = u'\u0653'
+HAMZA_ABOVE = u'\u0654'
+HAMZA_BELOW = u'\u0655'
+ZERO = u'\u0660'
+ONE = u'\u0661'
+TWO = u'\u0662'
+THREE = u'\u0663'
+FOUR = u'\u0664'
+FIVE = u'\u0665'
+SIX = u'\u0666'
+SEVEN = u'\u0667'
+EIGHT = u'\u0668'
+NINE = u'\u0669'
+PERCENT = u'\u066a'
+DECIMAL = u'\u066b'
+THOUSANDS = u'\u066c'
+STAR = u'\u066d'
+MINI_ALEF = u'\u0670'
+ALEF_WASLA = u'\u0671'
+FULL_STOP = u'\u06d4'
+BYTE_ORDER_MARK = u'\ufeff'
+
+# Diacritics
+FATHATAN = u'\u064b'
+DAMMATAN = u'\u064c'
+KASRATAN = u'\u064d'
+FATHA = u'\u064e'
+DAMMA = u'\u064f'
+KASRA = u'\u0650'
+SHADDA = u'\u0651'
+SUKUN = u'\u0652'
+
+#Ligatures
+LAM_ALEF = u'\ufefb'
+LAM_ALEF_HAMZA_ABOVE = u'\ufef7'
+LAM_ALEF_HAMZA_BELOW = u'\ufef9'
+LAM_ALEF_MADDA_ABOVE = u'\ufef5'
+SIMPLE_LAM_ALEF = u'\u0644\u0627'
+SIMPLE_LAM_ALEF_HAMZA_ABOVE = u'\u0644\u0623'
+SIMPLE_LAM_ALEF_HAMZA_BELOW = u'\u0644\u0625'
+SIMPLE_LAM_ALEF_MADDA_ABOVE = u'\u0644\u0622'
+
+
+HARAKAT_PAT = re.compile(u"["+u"".join([FATHATAN, DAMMATAN, KASRATAN,
+                                        FATHA, DAMMA, KASRA, SUKUN,
+                                        SHADDA])+u"]")
+HAMZAT_PAT = re.compile(u"["+u"".join([WAW_HAMZA, YEH_HAMZA])+u"]")
+ALEFAT_PAT = re.compile(u"["+u"".join([ALEF_MADDA, ALEF_HAMZA_ABOVE,
+                                       ALEF_HAMZA_BELOW, HAMZA_ABOVE,
+                                       HAMZA_BELOW])+u"]")
+LAMALEFAT_PAT = re.compile(u"["+u"".join([LAM_ALEF,
+                                          LAM_ALEF_HAMZA_ABOVE,
+                                          LAM_ALEF_HAMZA_BELOW,
+LAM_ALEF_MADDA_ABOVE])+u"]")
+
+def strip_tashkeel(text):
+    text = HARAKAT_PAT.sub('', text)
+    text = re.sub(u"[\u064E]", "", text,  flags=re.UNICODE) # fattha
+    text = re.sub(u"[\u0671]", "", text,  flags=re.UNICODE) # waSla
+    return text
+
+def strip_tatweel(text):
+    return re.sub(u'[%s]' % TATWEEL, '', text)
+
+# remove removing Tashkeel + removing Tatweel + non Arabic chars
+def remove_non_arabic(text):
+    text = strip_tashkeel(text)
+    text = strip_tatweel(text)
+    return ' '.join(re.sub(u"[^\u0621-\u063A\u0641-\u064A ]", " ", text,  flags=re.UNICODE).split())
+

spaces/arabic-auto-tashkeel/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

spaces/arabic-auto-tashkeel/README.md

@@ -0,0 +1,13 @@
+---
+title: Arabic Auto Tashkeel
+emoji: 🔥
+colorFrom: indigo
+colorTo: yellow
+sdk: gradio
+sdk_version: 4.27.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

spaces/arabic-auto-tashkeel/app.py

@@ -0,0 +1,104 @@
+import spaces
+import gradio as gr
+from shakkala import Shakkala
+from pathlib import Path
+import torch
+from eo_pl import TashkeelModel as TashkeelModelEO
+from ed_pl import TashkeelModel as TashkeelModelED
+from tashkeel_tokenizer import TashkeelTokenizer
+from utils import remove_non_arabic
+
+# Initialize the Shakkala model
+sh = Shakkala(version=3)
+model, graph = sh.get_model()
+
+def infer_shakkala(input_text):
+    input_int = sh.prepare_input(input_text)
+    logits = model.predict(input_int)[0]
+    predicted_harakat = sh.logits_to_text(logits)
+    final_output = sh.get_final_text(input_text, predicted_harakat)
+    print(final_output)
+    return final_output
+
+# Initialize the CaTT model and tokenizer
+tokenizer = TashkeelTokenizer()
+eo_ckpt_path = Path(__file__).parent / 'models/best_eo_mlm_ns_epoch_193.pt'
+ed_ckpt_path = Path(__file__).parent / 'models/best_ed_mlm_ns_epoch_178.pt'
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+max_seq_len = 1024
+print('Creating Model...')
+eo_model = TashkeelModelEO(tokenizer, max_seq_len=max_seq_len, n_layers=6, learnable_pos_emb=False)
+ed_model = TashkeelModelED(tokenizer, max_seq_len=max_seq_len, n_layers=3, learnable_pos_emb=False)
+
+eo_model.load_state_dict(torch.load(eo_ckpt_path, map_location=device))
+eo_model.eval().to(device)
+ed_model.load_state_dict(torch.load(ed_ckpt_path, map_location=device))
+ed_model.eval().to(device)
+
+@spaces.GPU()
+def infer_catt(input_text, choose_model):
+    input_text = remove_non_arabic(input_text)
+    batch_size = 16
+    verbose = True
+    if choose_model == 'Encoder-Only':
+        output_text = eo_model.do_tashkeel_batch([input_text], batch_size, verbose)
+    else:
+        output_text = ed_model.do_tashkeel_batch([input_text], batch_size, verbose)
+
+    return output_text[0]
+
+
+
+examples = ["السلام عليكم ورحمة الله وبركاته", "العلم نور", "الحمد لله"]
+
+with gr.Blocks(title="Arabic Tashkeel") as demo:
+    gr.HTML("<center><h1>Arabic Tashkeel</h1></center>")
+    gr.HTML(
+        "<center><p>Compare different methods for adding tashkeel to Arabic text.</p></center>"
+    )
+    with gr.Tab(label="CATT"):
+        gr.HTML("<center><h2>CATT: Character-based Arabic Tashkeel Transformer</h2></center>")
+        gr.HTML("<center><a href='https://github.com/abjadai/catt'>GitHub</a> - <a href='https://arxiv.org/abs/2407.03236'>Arxiv Paper</a></center>")
+
+        with gr.Row():
+            with gr.Column():
+                text_input1 = gr.Textbox(label="Input Text", rtl=True, text_align="right")
+                choose_model = gr.Radio(
+                    label="Choose Model",
+                    value="Encoder-Decoder",
+                    choices=["Encoder-Only", "Encoder-Decoder"],
+                )
+                with gr.Row():
+                    clear_button1 = gr.Button(value="Clear", variant="secondary")
+                    submit_button1 = gr.Button(value="Add Tashkeel", variant="primary")
+
+            with gr.Column():
+                text_output1 = gr.Textbox(label="Output Text", rtl=True, text_align="right")
+
+        gr.Examples(examples, text_input1, cache_examples=False)
+
+        submit_button1.click(infer_catt, inputs=[text_input1, choose_model], outputs=text_output1)
+        clear_button1.click(lambda: ("", ""), outputs=[text_input1, text_output1])
+
+    with gr.Tab(label="Shakkala"):
+        gr.HTML("<center><h2>Shakkala: Arabic Diacritization</h2></center>")
+        gr.HTML("<center><a href='https://github.com/Barqawiz/Shakkala'>GitHub</a> - <a href='https://pypi.org/project/shakkala/'>PyPi Package</a></center>")
+        with gr.Row():
+            with gr.Column():
+                text_input2 = gr.Textbox(label="Input Text", rtl=True, text_align="right")
+                with gr.Row():
+                    clear_button2 = gr.Button(value="Clear", variant="secondary")
+                    submit_button2 = gr.Button(value="Apply Tashkeel", variant="primary")
+            with gr.Column():
+                text_output2 = gr.Textbox(
+                    lines=1, label="Output Text", rtl=True, text_align="right"
+                )
+
+        submit_button2.click(infer_shakkala, inputs=text_input2, outputs=text_output2)
+        clear_button2.click(lambda: ("", ""), outputs=[text_input2, text_output2])
+        
+        gr.Examples(examples, text_input2, cache_examples=False)
+
+if __name__ == '__main__':
+    demo.queue().launch()

spaces/arabic-auto-tashkeel/bw2ar.py

@@ -0,0 +1,126 @@
+"""
+functions to convert Arabic words/text into buckwalter encoding and vice versa
+"""
+
+import sys
+import re
+import utils
+
+buck2uni = {
+            "'": u"\u0621", # hamza-on-the-line
+            "|": u"\u0622", # madda
+            ">": u"\u0623", # hamza-on-'alif
+            "&": u"\u0624", # hamza-on-waaw
+            "<": u"\u0625", # hamza-under-'alif
+            "}": u"\u0626", # hamza-on-yaa'
+            "A": u"\u0627", # bare 'alif
+            "b": u"\u0628", # baa'
+            "p": u"\u0629", # taa' marbuuTa
+            "t": u"\u062A", # taa'
+            "v": u"\u062B", # thaa'
+            "j": u"\u062C", # jiim
+            "H": u"\u062D", # Haa'
+            "x": u"\u062E", # khaa'
+            "d": u"\u062F", # daal
+            "*": u"\u0630", # dhaal
+            "r": u"\u0631", # raa'
+            "z": u"\u0632", # zaay
+            "s": u"\u0633", # siin
+            "$": u"\u0634", # shiin
+            "S": u"\u0635", # Saad
+            "D": u"\u0636", # Daad
+            "T": u"\u0637", # Taa'
+            "Z": u"\u0638", # Zaa' (DHaa')
+            "E": u"\u0639", # cayn
+            "g": u"\u063A", # ghayn
+            "_": u"\u0640", # taTwiil
+            "f": u"\u0641", # faa'
+            "q": u"\u0642", # qaaf
+            "k": u"\u0643", # kaaf
+            "l": u"\u0644", # laam
+            "m": u"\u0645", # miim
+            "n": u"\u0646", # nuun
+            "h": u"\u0647", # haa'
+            "w": u"\u0648", # waaw
+            "Y": u"\u0649", # 'alif maqSuura
+            "y": u"\u064A", # yaa'
+            "F": u"\u064B", # fatHatayn
+            "N": u"\u064C", # Dammatayn
+            "K": u"\u064D", # kasratayn
+            "a": u"\u064E", # fatHa
+            "u": u"\u064F", # Damma
+            "i": u"\u0650", # kasra
+            "~": u"\u0651", # shaddah
+            "o": u"\u0652", # sukuun
+            "`": u"\u0670", # dagger 'alif
+            "{": u"\u0671", # waSla
+}
+
+# For a reverse transliteration (Unicode -> Buckwalter), a dictionary
+# which is the reverse of the above buck2uni is essential.
+uni2buck = {}
+
+# Iterate through all the items in the buck2uni dict.
+for (key, value) in buck2uni.items():
+    # The value from buck2uni becomes a key in uni2buck, and vice
+    # versa for the keys.
+    uni2buck[value] = key
+
+# add special characters
+uni2buck[u"\ufefb"] = "lA"
+uni2buck[u"\ufef7"] = "l>"
+uni2buck[u"\ufef5"] = "l|"
+uni2buck[u"\ufef9"] = "l<"
+
+# clean the arabic text from unwanted characters that may cause problem while building the language model
+def clean_text(text):
+    text = re.sub(u"[\ufeff]", "", text,  flags=re.UNICODE) # strip Unicode Character 'ZERO WIDTH NO-BREAK SPACE' (U+FEFF). For more info, check http://www.fileformat.info/info/unicode/char/feff/index.htm
+    text = utils.remove_non_arabic(text)
+    text = utils.strip_tashkeel(text)
+    text = utils.strip_tatweel(text)
+    return text
+
+# convert a single word into buckwalter and vice versa
+def transliterate_word(input_word, direction='bw2ar'):
+    output_word = ''
+    # Loop over each character in the string, bw_word.
+    for char in input_word:
+        # Look up current char in the dictionary to get its
+        # respective value. If there is no match, e.g., chars like
+        # spaces, then just stick with the current char without any
+        # conversion.
+        # if type(char) == bytes:
+        #    char = char.decode('ascii')
+        if direction == 'bw2ar':
+            #print('in bw2ar')
+            output_word += buck2uni.get(char, char)
+        elif direction == 'ar2bw':
+            #print('in ar2bw')
+            output_word += uni2buck.get(char, char)
+        else:
+            sys.stderr.write('Error: invalid direction!')
+            sys.exit()
+    return output_word
+
+
+# convert a text into buckwalter and vice versa
+def transliterate_text(input_text, direction='bw2ar'):
+    output_text = ''
+    for input_word in input_text.split(' '):
+        output_text += transliterate_word(input_word, direction) + ' '
+
+    return output_text[:-1] # remove the last space ONLY
+
+
+if __name__ == '__main__':
+    if len(sys.argv) < 2:
+        sys.stderr.write('Usage: INPUT TEXT | python {} DIRECTION(bw2ar|ar2bw)'.format(sys.argv[1]))
+        exit(1)
+    for line in sys.stdin:
+        line = line if sys.argv[1] == 'bw2ar' else clean_text(line)
+        output_text = transliterate_text(line, direction=str(sys.argv[1]))
+        if output_text.strip() != '':
+            sys.stdout.write('{}\n'.format(output_text.strip()))
+
+
+

spaces/arabic-auto-tashkeel/ed.py

@@ -0,0 +1,63 @@
+
+import torch.nn as nn
+from transformer import *
+
+class Transformer(nn.Module):
+
+    def __init__(self, src_pad_idx, trg_pad_idx, enc_voc_size, dec_voc_size, d_model, n_head, max_len,
+                 ffn_hidden, n_layers, drop_prob, learnable_pos_emb=True):
+        super().__init__()
+        self.src_pad_idx = src_pad_idx
+        self.trg_pad_idx = trg_pad_idx
+        self.encoder = Encoder(d_model=d_model,
+                               n_head=n_head,
+                               max_len=max_len,
+                               ffn_hidden=ffn_hidden,
+                               enc_voc_size=enc_voc_size,
+                               drop_prob=drop_prob,
+                               n_layers=n_layers,
+                               padding_idx=src_pad_idx,
+                               learnable_pos_emb=learnable_pos_emb)
+
+        self.decoder = Decoder(d_model=d_model,
+                               n_head=n_head,
+                               max_len=max_len,
+                               ffn_hidden=ffn_hidden,
+                               dec_voc_size=dec_voc_size,
+                               drop_prob=drop_prob,
+                               n_layers=n_layers,
+                               padding_idx=trg_pad_idx,
+                               learnable_pos_emb=learnable_pos_emb)
+
+    def get_device(self):
+        return next(self.parameters()).device
+
+    def forward(self, src, trg):
+        device = self.get_device()
+        src_mask = self.make_pad_mask(src, src, self.src_pad_idx, self.src_pad_idx).to(device)
+        src_trg_mask = self.make_pad_mask(trg, src, self.trg_pad_idx, self.src_pad_idx).to(device)
+        trg_mask = self.make_pad_mask(trg, trg, self.trg_pad_idx, self.trg_pad_idx).to(device) * \
+                   self.make_no_peak_mask(trg, trg).to(device)
+        enc_src = self.encoder(src, src_mask)
+        output = self.decoder(trg, enc_src, trg_mask, src_trg_mask)
+        return output
+
+    def make_pad_mask(self, q, k, q_pad_idx, k_pad_idx):
+        len_q, len_k = q.size(1), k.size(1)
+        # batch_size x 1 x 1 x len_k
+        k = k.ne(k_pad_idx).unsqueeze(1).unsqueeze(2)
+        # batch_size x 1 x len_q x len_k
+        k = k.repeat(1, 1, len_q, 1)
+        # batch_size x 1 x len_q x 1
+        q = q.ne(q_pad_idx).unsqueeze(1).unsqueeze(3)
+        # batch_size x 1 x len_q x len_k
+        q = q.repeat(1, 1, 1, len_k)
+        mask = k & q
+        return mask
+
+    def make_no_peak_mask(self, q, k):
+        len_q, len_k = q.size(1), k.size(1)
+        # len_q x len_k
+        mask = torch.tril(torch.ones(len_q, len_k)).type(torch.BoolTensor)
+        return mask
+

spaces/arabic-auto-tashkeel/ed_pl.py

@@ -0,0 +1,164 @@
+
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from ed import Transformer
+from tqdm import tqdm
+import math
+import torch
+import torch.nn as nn
+
+from torch.nn.utils.rnn import pad_sequence
+# sequences is a list of tensors of shape TxH where T is the seqlen and H is the feats dim
+def pad_seq(sequences, batch_first=True, padding_value=0.0, prepadding=True):
+    lens = [i.shape[0]for i in sequences]
+    padded_sequences = pad_sequence(sequences, batch_first=True, padding_value=padding_value) # NxTxH
+    if prepadding:
+        for i in range(len(lens)):
+            padded_sequences[i] = padded_sequences[i].roll(-lens[i])
+    if not batch_first:
+        padded_sequences = padded_sequences.transpose(0, 1) # TxNxH
+    return padded_sequences
+
+
+
+def get_batches(X, batch_size=16):
+    num_batches = math.ceil(len(X) / batch_size)
+    for i in range(num_batches):
+        x = X[i*batch_size : (i+1)*batch_size]
+        yield x
+
+
+class TashkeelModel(pl.LightningModule):
+    def __init__(self, tokenizer, max_seq_len, d_model=512, n_layers=3, n_heads=16, drop_prob=0.1, learnable_pos_emb=True):
+
+        super(TashkeelModel, self).__init__()
+
+        ffn_hidden = 4 * d_model
+        src_pad_idx = tokenizer.letters_map['<PAD>']
+        trg_pad_idx = tokenizer.tashkeel_map['<PAD>']
+        enc_voc_size = len(tokenizer.letters_map) # 37 + 3
+        dec_voc_size = len(tokenizer.tashkeel_map) # 15 + 3
+        self.transformer = Transformer(src_pad_idx=src_pad_idx,
+                            trg_pad_idx=trg_pad_idx,
+                            d_model=d_model,
+                            enc_voc_size=enc_voc_size,
+                            dec_voc_size=dec_voc_size,
+                            max_len=max_seq_len,
+                            ffn_hidden=ffn_hidden,
+                            n_head=n_heads,
+                            n_layers=n_layers,
+                            drop_prob=drop_prob,
+                            learnable_pos_emb=learnable_pos_emb
+                            )
+
+        self.criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.tashkeel_map['<PAD>'])
+        self.tokenizer = tokenizer
+
+
+    def forward(self, x, y=None):
+        y_pred = self.transformer(x, y)
+        return y_pred
+
+
+    def training_step(self, batch, batch_idx):
+        input_ids, target_ids = batch
+        input_ids = input_ids[:, :-1]
+        y_in = target_ids[:, :-1]
+        y_out = target_ids[:, 1:]
+        y_pred = self(input_ids, y_in)
+        loss = self.criterion(y_pred.transpose(1, 2), y_out)
+
+        self.log('train_loss', loss, prog_bar=True)
+        sch = self.lr_schedulers()
+        sch.step()
+        self.log('lr', sch.get_last_lr()[0], prog_bar=True)
+        return loss
+
+
+    def validation_step(self, batch, batch_idx):
+        input_ids, target_ids = batch
+        input_ids = input_ids[:, :-1]
+        y_in = target_ids[:, :-1]
+        y_out = target_ids[:, 1:]
+        y_pred = self(input_ids, y_in)
+        loss = self.criterion(y_pred.transpose(1, 2), y_out)
+
+        pred_text_with_tashkeels = self.tokenizer.decode(input_ids, y_pred.argmax(2).squeeze())
+        true_text_with_tashkeels = self.tokenizer.decode(input_ids, y_out)
+        total_val_der_distance = 0
+        total_val_der_ref_length = 0
+        for i in range(len(true_text_with_tashkeels)):
+            pred_text_with_tashkeel = pred_text_with_tashkeels[i]
+            true_text_with_tashkeel = true_text_with_tashkeels[i]
+            val_der = self.tokenizer.compute_der(true_text_with_tashkeel, pred_text_with_tashkeel)
+            total_val_der_distance += val_der['distance']
+            total_val_der_ref_length += val_der['ref_length']
+
+        total_der_error = total_val_der_distance / total_val_der_ref_length
+        self.log('val_loss', loss)
+        self.log('val_der', torch.FloatTensor([total_der_error]))
+        self.log('val_der_distance', torch.FloatTensor([total_val_der_distance]))
+        self.log('val_der_ref_length', torch.FloatTensor([total_val_der_ref_length]))
+
+
+    def test_step(self, batch, batch_idx):
+        input_ids, target_ids = batch
+        y_pred = self(input_ids, None)
+        loss = self.criterion(y_pred.transpose(1, 2), target_ids)
+        self.log('test_loss', loss)
+
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=3e-4)
+        #max_iters = 10000
+        #lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_iters, eta_min=3e-6)
+        gamma = 1 / 1.000001
+        #gamma = 1 / 1.0001
+        lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma)
+        opts = {"optimizer": optimizer,  "lr_scheduler": lr_scheduler}
+        return opts
+
+
+    @torch.no_grad()
+    def do_tashkeel_batch(self, texts, batch_size=16, verbose=True):
+        self.eval()
+        device = next(self.parameters()).device
+        text_with_tashkeel = []
+        data_iter = get_batches(texts, batch_size)
+        if verbose:
+            num_batches = math.ceil(len(texts) / batch_size)
+            data_iter = tqdm(data_iter, total=num_batches)
+        for texts_mini in data_iter:
+            input_ids_list = []
+            for text in texts_mini:
+                input_ids, _ = self.tokenizer.encode(text, test_match=False)
+                input_ids_list.append(input_ids)
+            batch_input_ids = pad_seq(input_ids_list, batch_first=True, padding_value=self.tokenizer.letters_map['<PAD>'], prepadding=False)
+            target_ids = torch.LongTensor([[self.tokenizer.tashkeel_map['<BOS>']]] * len(texts_mini)).to(device)
+            src = batch_input_ids.to(device)
+
+            src_mask = self.transformer.make_pad_mask(src, src, self.transformer.src_pad_idx, self.transformer.src_pad_idx).to(device)
+            enc_src = self.transformer.encoder(src, src_mask)
+
+            for i in range(src.shape[1] - 1):
+                trg = target_ids
+                src_trg_mask = self.transformer.make_pad_mask(trg, src, self.transformer.trg_pad_idx, self.transformer.src_pad_idx).to(device)
+                trg_mask = self.transformer.make_pad_mask(trg, trg, self.transformer.trg_pad_idx, self.transformer.trg_pad_idx).to(device) * \
+                           self.transformer.make_no_peak_mask(trg, trg).to(device)
+
+                preds = self.transformer.decoder(trg, enc_src, trg_mask, src_trg_mask)
+                # IMPORTANT NOTE: the following code snippet is to FORCE the prediction of the input space char to output no_tashkeel tag '<NT>'
+                target_ids = torch.cat([target_ids, preds[:, -1].argmax(1).unsqueeze(1)], axis=1)
+                target_ids[self.tokenizer.letters_map[' '] == src[:, :target_ids.shape[1]]] = self.tokenizer.tashkeel_map[self.tokenizer.no_tashkeel_tag]
+#                target_ids = torch.cat([target_ids, preds[:, -1].argmax(1).unsqueeze(1)], axis=1)
+            text_with_tashkeel_mini = self.tokenizer.decode(src, target_ids)
+            text_with_tashkeel += text_with_tashkeel_mini
+        return text_with_tashkeel
+
+
+    @torch.no_grad()
+    def do_tashkeel(self, text):
+        return self.do_tashkeel_batch([text])[0]

spaces/arabic-auto-tashkeel/eo.py

@@ -0,0 +1,50 @@
+import torch.nn as nn
+from transformer import *
+
+class Transformer(nn.Module):
+    def __init__(self, src_pad_idx, trg_pad_idx, enc_voc_size, dec_voc_size, d_model, n_head, max_len,
+                 ffn_hidden, n_layers, drop_prob, learnable_pos_emb=True):
+        super().__init__()
+        self.src_pad_idx = src_pad_idx
+        self.trg_pad_idx = trg_pad_idx
+        self.encoder = Encoder(d_model=d_model,
+                               n_head=n_head,
+                               max_len=max_len,
+                               ffn_hidden=ffn_hidden,
+                               enc_voc_size=enc_voc_size,
+                               drop_prob=drop_prob,
+                               n_layers=n_layers,
+                               padding_idx=src_pad_idx,
+                               learnable_pos_emb=learnable_pos_emb)
+
+        self.decoder = nn.Linear(d_model, dec_voc_size)
+
+    def get_device(self):
+        return next(self.parameters()).device
+
+    def forward(self, src):
+        device = self.get_device()
+        src_mask = self.make_pad_mask(src, src, self.src_pad_idx, self.src_pad_idx).to(device)
+        enc_src = self.encoder(src, src_mask)
+        output = self.decoder(enc_src)
+        return output
+
+    def make_pad_mask(self, q, k, q_pad_idx, k_pad_idx):
+        len_q, len_k = q.size(1), k.size(1)
+        # batch_size x 1 x 1 x len_k
+        k = k.ne(k_pad_idx).unsqueeze(1).unsqueeze(2)
+        # batch_size x 1 x len_q x len_k
+        k = k.repeat(1, 1, len_q, 1)
+        # batch_size x 1 x len_q x 1
+        q = q.ne(q_pad_idx).unsqueeze(1).unsqueeze(3)
+        # batch_size x 1 x len_q x len_k
+        q = q.repeat(1, 1, 1, len_k)
+        mask = k & q
+        return mask
+
+    def make_no_peak_mask(self, q, k):
+        len_q, len_k = q.size(1), k.size(1)
+        # len_q x len_k
+        mask = torch.tril(torch.ones(len_q, len_k)).type(torch.BoolTensor)
+        return mask
+

spaces/arabic-auto-tashkeel/eo_pl.py

@@ -0,0 +1,151 @@
+
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from eo import Transformer
+from tqdm import tqdm
+import math
+import torch
+import torch.nn as nn
+
+from torch.nn.utils.rnn import pad_sequence
+# sequences is a list of tensors of shape TxH where T is the seqlen and H is the feats dim
+def pad_seq(sequences, batch_first=True, padding_value=0.0, prepadding=True):
+    lens = [i.shape[0]for i in sequences]
+    padded_sequences = pad_sequence(sequences, batch_first=True, padding_value=padding_value) # NxTxH
+    if prepadding:
+        for i in range(len(lens)):
+            padded_sequences[i] = padded_sequences[i].roll(-lens[i])
+    if not batch_first:
+        padded_sequences = padded_sequences.transpose(0, 1) # TxNxH
+    return padded_sequences
+
+
+
+def get_batches(X, batch_size=16):
+    num_batches = math.ceil(len(X) / batch_size)
+    for i in range(num_batches):
+        x = X[i*batch_size : (i+1)*batch_size]
+        yield x
+
+
+class TashkeelModel(pl.LightningModule):
+    def __init__(self, tokenizer, max_seq_len, d_model=512, n_layers=3, n_heads=16, drop_prob=0.1, learnable_pos_emb=True):
+
+        super(TashkeelModel, self).__init__()
+
+        ffn_hidden = 4 * d_model
+        src_pad_idx = tokenizer.letters_map['<PAD>']
+        trg_pad_idx = tokenizer.tashkeel_map['<PAD>']
+        enc_voc_size = len(tokenizer.letters_map) # 37 + 3
+        dec_voc_size = len(tokenizer.tashkeel_map) # 15 + 3
+        self.transformer = Transformer(src_pad_idx=src_pad_idx,
+                            trg_pad_idx=trg_pad_idx,
+                            d_model=d_model,
+                            enc_voc_size=enc_voc_size,
+                            dec_voc_size=dec_voc_size,
+                            max_len=max_seq_len,
+                            ffn_hidden=ffn_hidden,
+                            n_head=n_heads,
+                            n_layers=n_layers,
+                            drop_prob=drop_prob,
+                            learnable_pos_emb=learnable_pos_emb
+                            )
+
+        self.criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.tashkeel_map['<PAD>'])
+        self.tokenizer = tokenizer
+
+
+    def forward(self, x):
+        y_pred = self.transformer(x)
+        return y_pred
+
+
+    def training_step(self, batch, batch_idx):
+        input_ids, target_ids = batch
+        input_ids = input_ids[:, 1:-1]
+        y_out = target_ids[:, 1:-1]
+        y_pred = self(input_ids)
+        loss = self.criterion(y_pred.transpose(1, 2), y_out)
+
+        self.log('train_loss', loss, prog_bar=True)
+#        sch = self.lr_schedulers()
+#        sch.step()
+#        self.log('lr', sch.get_last_lr()[0], prog_bar=True)
+        return loss
+
+
+    def validation_step(self, batch, batch_idx):
+        input_ids, target_ids = batch
+        input_ids = input_ids[:, 1:-1]
+        y_out = target_ids[:, 1:-1]
+        y_pred = self(input_ids)
+        loss = self.criterion(y_pred.transpose(1, 2), y_out)
+
+        pred_text_with_tashkeels = self.tokenizer.decode(input_ids, y_pred.argmax(2).squeeze())
+        true_text_with_tashkeels = self.tokenizer.decode(input_ids, y_out)
+        total_val_der_distance = 0
+        total_val_der_ref_length = 0
+        for i in range(len(true_text_with_tashkeels)):
+            pred_text_with_tashkeel = pred_text_with_tashkeels[i]
+            true_text_with_tashkeel = true_text_with_tashkeels[i]
+            val_der = self.tokenizer.compute_der(true_text_with_tashkeel, pred_text_with_tashkeel)
+            total_val_der_distance += val_der['distance']
+            total_val_der_ref_length += val_der['ref_length']
+
+        total_der_error = total_val_der_distance / total_val_der_ref_length
+        self.log('val_loss', loss)
+        self.log('val_der', torch.FloatTensor([total_der_error]))
+        self.log('val_der_distance', torch.FloatTensor([total_val_der_distance]))
+        self.log('val_der_ref_length', torch.FloatTensor([total_val_der_ref_length]))
+
+
+    def test_step(self, batch, batch_idx):
+        input_ids, target_ids = batch
+        y_pred = self(input_ids, None)
+        loss = self.criterion(y_pred.transpose(1, 2), target_ids)
+        self.log('test_loss', loss)
+
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=3e-5)
+        #max_iters = 10000
+        #lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_iters, eta_min=3e-6)
+        gamma = 1 / 1.000001
+        #gamma = 1 / 1.0001
+        #lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma)
+        opts = {"optimizer": optimizer} #,  "lr_scheduler": lr_scheduler}
+        return opts
+
+
+    @torch.no_grad()
+    def do_tashkeel_batch(self, texts, batch_size=16, verbose=True):
+        self.eval()
+        device = next(self.parameters()).device
+        text_with_tashkeel = []
+        data_iter = get_batches(texts, batch_size)
+        if verbose:
+            num_batches = math.ceil(len(texts) / batch_size)
+            data_iter = tqdm(data_iter, total=num_batches)
+        for texts_mini in data_iter:
+            input_ids_list = []
+            for text in texts_mini:
+                input_ids, _ = self.tokenizer.encode(text, test_match=False)
+                input_ids_list.append(input_ids)
+            batch_input_ids = pad_seq(input_ids_list, batch_first=True, padding_value=self.tokenizer.letters_map['<PAD>'], prepadding=False)
+            batch_input_ids = batch_input_ids[:, 1:-1].to(device)
+            y_pred = self(batch_input_ids)
+            y_pred = y_pred.argmax(-1)
+            # IMPORTANT NOTE: the following code snippet is to FORCE the prediction of the input space char to output no_tashkeel tag '<NT>'
+            y_pred[self.tokenizer.letters_map[' '] == batch_input_ids] = self.tokenizer.tashkeel_map[self.tokenizer.no_tashkeel_tag]
+            text_with_tashkeel_mini = self.tokenizer.decode(batch_input_ids, y_pred)
+            text_with_tashkeel += text_with_tashkeel_mini
+
+        return text_with_tashkeel
+
+
+    @torch.no_grad()
+    def do_tashkeel(self, text):
+        return self.do_tashkeel_batch([text])[0]

spaces/arabic-auto-tashkeel/models/best_ed_mlm_ns_epoch_178.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6297ec54f9af518a59d4b9e5831d3ad4da304435b8c7fb969f6e3a16816ac23
+size 88403510