U4RASD/NeoAraBERT

1from typing import List, Tuple2from transformers import PreTrainedTokenizerFast3import re4import fast_disambig5 6_TATWEEL_RE = re.compile(r"\u0640")7_ALIF_RE = re.compile(r"[آأإٱ]")8_ALIF_MAK_RE = re.compile(r"ى")9_TEH_MARB_RE = re.compile(r"ة")10_ZERO_WIDTH_RE = re.compile(r"[\u200B-\u200D\u200E\u200F\uFEFF]")11ARABIC_DIACRITICS = {12    "ً", "ٌ", "ٍ",13    "َ", "ُ", "ِ",14    "ّ", "ْ",15    "ٗ", "٘", "ٙ", "ٚ", "ٛ", "ٜ", "ٝ", "ٞ", "ٟ",16    "ؐ", "ؑ", "ؒ", "ؓ", "ؔ", "ؕ", "ؖ", "ؗ", "ؘ", "ؙ", "ؚ",17    "ۖ", "ۗ", "ۘ", "ۙ", "ۚ", "ۛ", "ۜ", "۟", "۠", "ۡ", "ۢ", "ۣ", "ۤ", "ۧ", "ۨ",18    "۪", "۫", "۬", "ۭ",19}20 21def separate_diacritics(text):22    tokens = re.split(r'(\s+|\[\+\])', text)23    processed_tokens = []24 25    for token in tokens:26        if not token:27            continue28        if token.isspace() or token == '[+]':29            processed_tokens.append(token)30            continue31 32        if not any(c in ARABIC_DIACRITICS for c in token):33            processed_tokens.append(token)34            continue35 36        base_chars = []37        diac_groups = []38 39        for char in token:40            if char in ARABIC_DIACRITICS:41                if not diac_groups:42                    base_chars.append(" ")43                    diac_groups.append([])44                diac_groups[-1].append(char)45            else:46                base_chars.append(char)47                diac_groups.append([])48 49        base_word = "".join(base_chars)50        diac_string = []51        for group in diac_groups:52            if group:53                diac_string.append("".join(group))54            else:55                diac_string.append("◌")56 57        processed_tokens.append(base_word + " " + "".join(diac_string))58    return "".join(processed_tokens)59 60def normalize_arabic(text):61    text = _TATWEEL_RE.sub("", text)62    text = _ZERO_WIDTH_RE.sub("", text)63    text = _ALIF_RE.sub("ا", text)64    text = _ALIF_MAK_RE.sub("ي", text)65    text = _TEH_MARB_RE.sub("ه", text)66    return text67 68class ArabicMorphTokenizer(PreTrainedTokenizerFast):69    slow_tokenizer_class = None70 71    def __init__(self, tokenizer_file=None, apply_stemming=True, **kwargs):72        super().__init__(tokenizer_file=tokenizer_file, **kwargs)73        self.apply_stemming = apply_stemming74        if self.apply_stemming:75            self.stemmer = fast_disambig.camel.Stemmer()76 77 78    def _preprocess_one(self, s, do_stem):79        if isinstance(s, (list, tuple)):80            return [self._preprocess_one(x, do_stem) for x in s]81        if do_stem:82            s = self.stemmer.stem(s, preserve_diacritics=True)83        s = normalize_arabic(s)84        s = separate_diacritics(s)85        return s86 87    def _preprocess_pair(self, text, text_pair, do_stem):88        def maybe(s):89            return self._preprocess_one(s, do_stem) if isinstance(s, str) else s90        if isinstance(text, (list, tuple)):91            text = [maybe(x) for x in text]92        else:93            text = maybe(text)94        if isinstance(text_pair, (list, tuple)):95            text_pair = [maybe(x) for x in text_pair]96        else:97            text_pair = maybe(text_pair)98        return text, text_pair99 100    def _pop_flag(self, kwargs):101        v = kwargs.pop("apply_stemming", None)102        return self.apply_stemming if v is None else bool(v)103 104    def __call__(self, text=None, text_pair=None, *args, **kwargs):105        flag = self._pop_flag(kwargs)106        if not getattr(self, "_processing", False):107            self._processing = True108            try:109                text, text_pair = self._preprocess_pair(text, text_pair, flag)110                return super().__call__(text=text, text_pair=text_pair, *args, **kwargs)111            finally:112                self._processing = False113        return super().__call__(text=text, text_pair=text_pair, *args, **kwargs)114 115    def encode(self, text, text_pair=None, *args, **kwargs):116        flag = self._pop_flag(kwargs)117        if not getattr(self, "_processing", False):118            self._processing = True119            try:120                text, text_pair = self._preprocess_pair(text, text_pair, flag)121                return super().encode(text, text_pair, *args, **kwargs)122            finally:123                self._processing = False124        return super().encode(text, text_pair, *args, **kwargs)125 126    def encode_plus(self, text=None, text_pair=None, *args, **kwargs):127        flag = self._pop_flag(kwargs)128        if not getattr(self, "_processing", False):129            self._processing = True130            try:131                text, text_pair = self._preprocess_pair(text, text_pair, flag)132                return super().encode_plus(text=text, text_pair=text_pair, *args, **kwargs)133            finally:134                self._processing = False135        return super().encode_plus(text=text, text_pair=text_pair, *args, **kwargs)136 137    def batch_encode_plus(self, batch_text_or_text_pairs=None, *args, **kwargs):138        flag = self._pop_flag(kwargs)139        if not getattr(self, "_processing", False):140            self._processing = True141            try:142                data = batch_text_or_text_pairs143                if isinstance(data, (list, tuple)):144                    new_data = []145                    for item in data:146                        if isinstance(item, (list, tuple)) and len(item) == 2:147                            new_data.append(self._preprocess_pair(item[0], item[1], flag))148                        else:149                            new_data.append(self._preprocess_one(item, flag))150                    batch_text_or_text_pairs = new_data151                return super().batch_encode_plus(batch_text_or_text_pairs=batch_text_or_text_pairs, *args, **kwargs)152            finally:153                self._processing = False154        return super().batch_encode_plus(batch_text_or_text_pairs=batch_text_or_text_pairs, *args, **kwargs)155 156    def preprocess(self, text, apply_stemming=True):157        flag = self.apply_stemming if apply_stemming is None else bool(apply_stemming)158        return self._preprocess_one(text, flag)159