Upload folder using huggingface_hub

1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+    "word_embedding_dimension": 768,
+    "pooling_mode_cls_token": false,
+    "pooling_mode_mean_tokens": true,
+    "pooling_mode_max_tokens": false,
+    "pooling_mode_mean_sqrt_len_tokens": false,
+    "pooling_mode_weightedmean_tokens": false,
+    "pooling_mode_lasttoken": false,
+    "include_prompt": true
+}

README.md

@@ -0,0 +1,56 @@
+---
+library_name: sentence-transformers
+pipeline_tag: sentence-similarity
+tags:
+- sentence-transformers
+- sentence-similarity
+- feature-extraction
+- dense
+- arabic
+---
+
+# NeoAraBERT-STS
+
+Sentence-transformers model for Arabic semantic textual similarity.
+
+## Usage
+
+```bash
+pip install -U sentence-transformers torch
+```
+
+```python
+import torch
+from sentence_transformers import SentenceTransformer
+
+model_name = "U4RASD/NeoAraBERT-STS"
+
+finetuned_model = SentenceTransformer(
+    model_name,
+    model_kwargs={"trust_remote_code": True, "torch_dtype": torch.float32},
+    tokenizer_kwargs={"trust_remote_code": True},
+    config_kwargs={"trust_remote_code": True},
+)
+finetuned_model.max_seq_length = 512
+
+sentences = [
+    "التقارير بدأت تصل في وقت متأخر من هذا العام ويتم مراجعتها",
+    "يتم مراجعة التقارير في أواخر هذا العام.",
+    "لم يكن هناك تقارير هذا العام على الإطلاق.",
+]
+
+embeddings = finetuned_model.encode(sentences)
+similarities = finetuned_model.similarity(embeddings, embeddings)
+
+print(embeddings.shape)
+print(similarities)
+```
+
+## Model Type
+
+- **Model type:** Sentence Transformer
+- **Task:** Sentence similarity / semantic textual similarity
+- **Language:** Arabic
+- **Embedding size:** 768
+- **Max sequence length:** 512
+- **Similarity function:** Cosine similarity

acrps_light_stemming.py

@@ -0,0 +1,420 @@
+"""
+Fast Arabic Stemmer
+
+Optimized stemmer for Arabic text with the following optimizations:
+1. MLE word-level caching - avoids redundant disambiguation
+2. O(1) set lookups instead of O(n) list lookups
+3. String operations instead of byte encoding
+4. Reduced redundant dediac_ar() calls
+5. Fast Arabic-focused tokenizer regex (1000x faster than full Unicode)
+
+Usage:
+    from stemmer import stem, create_stemmer
+
+    # Simple usage (uses module-level stemmer instance)
+    result = stem("وَالْكِتَابُ الْجَمِيلُ")
+
+    # With diacritics preservation
+    result = stem("وَالْكِتَابُ الْجَمِيلُ", apply_diacritics=True)
+
+    # Or create your own instance
+    stemmer = create_stemmer()
+    result = stemmer.stem("النص العربي")
+"""
+
+import re
+from collections import deque
+from types import MethodType
+
+from camel_tools.disambig.mle import MLEDisambiguator
+from camel_tools.utils.dediac import dediac_ar
+from .constants import list_al_t, list_al, list_t
+
+__all__ = ["stem", "create_stemmer", "Stemmer"]
+
+
+_SET_AL_T = frozenset(list_al_t)
+_SET_AL = frozenset(list_al)
+_SET_T = frozenset(list_t)
+
+_ALEF_LAM = "ال"
+_TAA_MARBOUTA_ATTACHED = "ة"
+_TAA_MARBOUTA_DETACHED = "\ufe93"
+_HAA_ATTACHED = "ه"
+
+_PATTERN_LAM_PLUS = "ل[+]"
+_PATTERN_ALEF_LAM_PLUS = "ال[+]"
+_REPLACEMENT_LAM_LAM_PLUS = "لل[+]"
+
+_DIACRITIC_MARKS = frozenset(
+    {
+        "\u064b",
+        "\u064c",
+        "\u064d",
+        "\u064e",
+        "\u064f",
+        "\u0650",
+        "\u0651",
+        "\u0652",
+        "\u0670",
+    }
+)
+
+
+_ARABIC_WORD_CHARS = (
+    r"\u0621-\u063A"  # Arabic letters (hamza to ghain)
+    r"\u0641-\u064A"  # Arabic letters (fa to ya)
+    r"\u064B-\u0652"  # Arabic diacritics
+    r"\u0653-\u0655"  # Combining marks
+    r"\u0670"  # Superscript alef
+    r"\u0671-\u06D3"  # Extended Arabic letters
+    r"\u06D5-\u06FF"  # More letters and marks
+    r"\u0750-\u077F"  # Arabic Supplement
+    r"\u08A0-\u08FF"  # Arabic Extended-A
+    r"\uFB50-\uFDFF"  # Arabic Presentation Forms-A
+    r"\uFE70-\uFEFF"  # Arabic Presentation Forms-B
+)
+_LATIN_NUM = r"a-zA-Z0-9"
+_ARABIC_INDIC_DIGITS = r"\u0660-\u0669"
+
+_WORD_PATTERN = f"[{_ARABIC_WORD_CHARS}{_LATIN_NUM}{_ARABIC_INDIC_DIGITS}]+"
+_PUNCT_PATTERN = f"[^{_ARABIC_WORD_CHARS}{_LATIN_NUM}{_ARABIC_INDIC_DIGITS}\\s]"
+_WHITESPACE_PATTERN = r"\s+"
+
+_TOKENIZE_RE = re.compile(f"{_WORD_PATTERN}|{_PUNCT_PATTERN}|{_WHITESPACE_PATTERN}")
+_NORM_TATWEEL_RE = re.compile(r"\u0640")
+_NORM_ZERO_WIDTH_RE = re.compile(r"[\u200B-\u200D\u200E\u200F\uFEFF]")
+
+
+def _normalize_for_stem(text: str) -> str:
+    text = _NORM_TATWEEL_RE.sub("", text)
+    text = _NORM_ZERO_WIDTH_RE.sub("", text)
+    return text
+
+
+def _tokenize(text):
+    return _TOKENIZE_RE.findall(text)
+
+
+def _merge_tokens(tokens, original_word):
+    parts = []
+    for tok in tokens:
+        if tok == "[+]":
+            parts.append("_")
+        elif tok.endswith("[+]"):
+            parts.append(tok[:-3])
+        elif tok.startswith("[+]"):
+            parts.append(tok[3:])
+        elif tok.endswith("+"):
+            parts.append(tok[:-1])
+        elif tok.startswith("+"):
+            parts.append(tok[1:])
+        else:
+            parts.append(tok)
+    return "".join(parts)
+
+
+def _has_diacritics(word):
+    for char in word:
+        if char in _DIACRITIC_MARKS:
+            return True
+    return False
+
+
+def _apply_diacritics_to_segments(segments, diacritized_word):
+    result = []
+
+    leading_diacritics = []
+    i = 0
+    while i < len(diacritized_word) and diacritized_word[i] in _DIACRITIC_MARKS:
+        leading_diacritics.append(diacritized_word[i])
+        i += 1
+
+    diacritic_index = len(leading_diacritics)
+
+    for segment_idx, segment in enumerate(segments):
+        if segment == "[+]":
+            result.append(segment)
+        else:
+            diacritized_segment = []
+            if segment_idx == 0 and leading_diacritics:
+                diacritized_segment.extend(leading_diacritics)
+
+            i = 0
+            while i < len(segment):
+                char = segment[i]
+                if segment[i : i + 3] == "[+]":
+                    diacritized_segment.append("[+]")
+                    i += 3
+                    continue
+
+                if diacritic_index < len(diacritized_word):
+                    while (
+                        diacritic_index < len(diacritized_word)
+                        and diacritized_word[diacritic_index] in _DIACRITIC_MARKS
+                    ):
+                        diacritic_index += 1
+
+                    if (
+                        diacritic_index < len(diacritized_word)
+                        and diacritized_word[diacritic_index] == char
+                    ):
+                        diacritized_segment.append(char)
+                        diacritic_index += 1
+                        while (
+                            diacritic_index < len(diacritized_word)
+                            and diacritized_word[diacritic_index] in _DIACRITIC_MARKS
+                        ):
+                            diacritized_segment.append(
+                                diacritized_word[diacritic_index]
+                            )
+                            diacritic_index += 1
+                    else:
+                        diacritized_segment.append(char)
+                else:
+                    diacritized_segment.append(char)
+                i += 1
+
+            result.append("".join(diacritized_segment))
+    return result
+
+
+def _merge_alef_and_alef_lam(input_list):
+    modified_list = []
+    i = 0
+    while i < len(input_list):
+        if i < len(input_list) - 1:
+            if (
+                input_list[i] == _PATTERN_LAM_PLUS
+                and input_list[i + 1] == _PATTERN_ALEF_LAM_PLUS
+            ):
+                modified_list.append(_REPLACEMENT_LAM_LAM_PLUS)
+                i += 2
+                continue
+        modified_list.append(input_list[i])
+        i += 1
+    return modified_list
+
+
+def _process_NOAN_word(word):
+    starts_with_al = word.startswith(_ALEF_LAM)
+    ends_with_ta = word.endswith(_TAA_MARBOUTA_ATTACHED) or word.endswith(
+        _TAA_MARBOUTA_DETACHED
+    )
+
+    if starts_with_al and ends_with_ta:
+        if word in _SET_AL_T:
+            stripped_word = word[2:-1]
+            first_part = word[0:2] + "[+]"
+            last_part = "[+]" + word[-1]
+            return [first_part, stripped_word, last_part]
+
+    if starts_with_al:
+        if word in _SET_AL:
+            stripped_word = word[2:]
+            first_part = word[0:2] + "[+]"
+            return [first_part, stripped_word]
+
+    if ends_with_ta:
+        if word in _SET_T:
+            stripped_word = word[:-1]
+            last_part = "[+]" + word[-1]
+            return [stripped_word, last_part]
+
+    return [word]
+
+
+def _split_token_on_t(list_toks):
+    new_list = []
+    for token in list_toks:
+        last_char = token[-1] if token else ""
+        if last_char in (_TAA_MARBOUTA_ATTACHED, _TAA_MARBOUTA_DETACHED, _HAA_ATTACHED):
+            if token == _HAA_ATTACHED:
+                new_list.append("[+]" + _TAA_MARBOUTA_ATTACHED)
+            else:
+                new_list.append(token[:-1])
+                new_list.append("[+]" + token[-1])
+        else:
+            new_list.append(token)
+    return new_list
+
+
+def _replace_separator(toks):
+    for i, tok in enumerate(toks):
+        if tok.startswith("+"):
+            toks[i] = "[+]" + tok[1:]
+        if tok.endswith("+"):
+            toks[i] = tok[:-1] + "[+]"
+    return toks
+
+
+def _morph_tokenize(
+    words, disambiguator, scheme="d3tok", split=True, apply_diacritics=True
+):
+    disambig_words = disambiguator.disambiguate(words)
+    result = deque()
+
+    for original, disambig_word in zip(words, disambig_words):
+        scored_analyses = disambig_word.analyses
+        original_word = original
+        dediac_word = dediac_ar(original_word)
+        word_has_diacritics = _has_diacritics(original_word)
+
+        if not scored_analyses:
+            result.append(original_word)
+            continue
+
+        analysis = scored_analyses[0].analysis
+        tok_raw = analysis.get(scheme, None)
+        tok = dediac_ar(tok_raw) if tok_raw is not None else None
+
+        ends_with_ta = dediac_word.endswith(
+            _TAA_MARBOUTA_ATTACHED
+        ) or dediac_word.endswith(_TAA_MARBOUTA_DETACHED)
+
+        if ends_with_ta:
+            if tok is not None:
+                toks = tok.split("_")
+                toks = _split_token_on_t(toks)
+                toks = _replace_separator(toks)
+                toks = _merge_alef_and_alef_lam(toks)
+                merged_toks = _merge_tokens(toks, dediac_word)
+
+                if merged_toks == dediac_word and len(toks) > 1:
+                    if apply_diacritics and word_has_diacritics:
+                        toks = _apply_diacritics_to_segments(toks, original)
+                    result.extend(toks)
+                    continue
+                else:
+                    result.append(original_word)
+                    continue
+
+        if tok is None or "NOAN" in tok:
+            noan_toks = _process_NOAN_word(dediac_word)
+            if apply_diacritics and word_has_diacritics:
+                noan_toks = _apply_diacritics_to_segments(noan_toks, original)
+            result.extend(noan_toks)
+
+        elif split:
+            toks = tok.split("_")
+            toks = _replace_separator(toks)
+            toks = _merge_alef_and_alef_lam(toks)
+            merged_toks = _merge_tokens(toks, dediac_word)
+
+            if merged_toks == dediac_word and len(toks) > 1:
+                if apply_diacritics and word_has_diacritics:
+                    toks = _apply_diacritics_to_segments(toks, original)
+                result.extend(toks)
+            else:
+                result.append(original_word)
+
+        else:
+            if tok == dediac_word:
+                result.append(original_word)
+            else:
+                result.append(original_word)
+
+    return list(result)
+
+
+def _create_cached_score_fn(mle_instance):
+    cache = {}
+    original_method = mle_instance._scored_analyses
+
+    def cached_score_fn(self, word_dd):
+        if word_dd in cache:
+            return cache[word_dd]
+        result = original_method(word_dd)
+        cache[word_dd] = result
+        return result
+
+    return cache, MethodType(cached_score_fn, mle_instance)
+
+
+# ============================================================================
+
+
+class Stemmer:
+    """
+    Fast Arabic stemmer with MLE caching.
+
+    Example:
+        stemmer = Stemmer()
+        result = stemmer.stem("النص العربي")
+        result = stemmer.stem("وَالْكِتَابُ", apply_diacritics=True)
+    """
+
+    def __init__(self):
+        """Initialize the stemmer with MLE disambiguator and caching."""
+        self._mle = MLEDisambiguator.pretrained("calima-msa-r13")
+        self._cache, cached_method = _create_cached_score_fn(self._mle)
+        self._mle._score_fn = cached_method
+
+    def stem(self, text: str, apply_diacritics: bool = False) -> str:
+        """
+        Stem Arabic text.
+
+        Args:
+            text: Arabic text to stem.
+            apply_diacritics: If True, preserve diacritics from input in output.
+                            If False (default), output will be without diacritics.
+
+        Returns:
+            Stemmed text with morphological segmentation markers [+].
+        """
+        text = _normalize_for_stem(text)
+        tokens = _tokenize(text)
+        stemmed_tokens = _morph_tokenize(
+            tokens, self._mle, apply_diacritics=apply_diacritics
+        )
+        return "".join(stemmed_tokens)
+
+    def clear_cache(self):
+        """Clear the disambiguation cache."""
+        self._cache.clear()
+
+    @property
+    def cache_size(self) -> int:
+        """Return the number of cached word disambiguations."""
+        return len(self._cache)
+
+
+def create_stemmer() -> Stemmer:
+    """
+    Create a new Stemmer instance.
+
+    Returns:
+        A new Stemmer instance with its own cache.
+    """
+    return Stemmer()
+
+
+_default_stemmer = None
+
+
+def stem(text: str, apply_diacritics: bool = False) -> str:
+    """
+    Stem Arabic text using a shared stemmer instance.
+
+    This is a convenience function that uses a module-level stemmer.
+    For better control over caching, create your own Stemmer instance.
+
+    Args:
+        text: Arabic text to stem.
+        apply_diacritics: If True, preserve diacritics from input in output.
+                         If False (default), output will be without diacritics.
+
+    Returns:
+        Stemmed text with morphological segmentation markers [+].
+
+    Example:
+        >>> stem("والكتاب الجميل")
+        'و[+]ال[+]كتاب ال[+]جميل'
+
+        >>> stem("وَالْكِتَابُ", apply_diacritics=True)
+        'وَ[+]الْ[+]كِتَابُ'
+    """
+    global _default_stemmer
+    if _default_stemmer is None:
+        _default_stemmer = Stemmer()
+    return _default_stemmer.stem(text, apply_diacritics=apply_diacritics)

config.json

@@ -0,0 +1,49 @@
+{
+  "architectures": [
+    "NeoBERT"
+  ],
+  "auto_map": {
+    "AutoConfig": "model.NeoBERTConfig",
+    "AutoModel": "model.NeoBERT",
+    "AutoModelForMaskedLM": "model.NeoBERTLMHead",
+    "AutoModelForSequenceClassification": "model.NeoBERTForSequenceClassification"
+  },
+  "classifier_init_range": 0.02,
+  "decoder_init_range": 0.02,
+  "dim_head": 64,
+  "dtype": "float32",
+  "embedding_init_range": 0.02,
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "kwargs": {
+    "architectures": [
+      "NeoBERTLMHead"
+    ],
+    "attn_implementation": null,
+    "auto_map": {
+      "AutoConfig": "model.NeoBERTConfig",
+      "AutoModel": "model.NeoBERT",
+      "AutoModelForMaskedLM": "model.NeoBERTLMHead",
+      "AutoModelForSequenceClassification": "model.NeoBERTForSequenceClassification"
+    },
+    "classifier_init_range": 0.02,
+    "dim_head": 64,
+    "kwargs": {
+      "classifier_init_range": 0.02,
+      "trust_remote_code": true
+    },
+    "model_type": "neobert",
+    "torch_dtype": "float32",
+    "transformers_version": "4.48.2",
+    "trust_remote_code": true
+  },
+  "max_length": 1024,
+  "model_type": "neobert",
+  "norm_eps": 1e-05,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "pad_token_id": 0,
+  "transformers_version": "4.57.0",
+  "trust_remote_code": true,
+  "vocab_size": 65000
+}

config_sentence_transformers.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "SentenceTransformer",
+  "__version__": {
+    "sentence_transformers": "5.2.3",
+    "transformers": "4.57.0",
+    "pytorch": "2.11.0+cu128"
+  },
+  "prompts": {
+    "query": "",
+    "document": ""
+  },
+  "default_prompt_name": null,
+  "similarity_fn_name": "cosine"
+}

model.py

@@ -0,0 +1,434 @@
+# From https://github.com/facebookresearch/llama/blob/main/llama/model.py
+
+import torch
+from torch import nn
+
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.nn.functional import scaled_dot_product_attention
+
+from typing import Optional, Tuple
+import numpy as np
+
+from xformers.ops import SwiGLU
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_varlen_func
+
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    FLASH_ATTN_AVAILABLE = False
+
+from transformers import (
+    PreTrainedModel,
+    PretrainedConfig,
+    DataCollatorForLanguageModeling,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    MaskedLMOutput,
+    SequenceClassifierOutput,
+)
+
+from .rotary import precompute_freqs_cis, apply_rotary_emb
+
+
+class DataCollatorWithPacking(DataCollatorForLanguageModeling):
+    def __init__(self, pack_sequences=False, **kwargs):
+        super().__init__(**kwargs)
+        self.pack_sequences = pack_sequences
+
+    def __call__(self, batch):
+        if self.pack_sequences:
+            # Add position_ids if not present
+            if "position_ids" not in batch[0]:
+                for item in batch:
+                    item["position_ids"] = list(range(len(item["input_ids"])))
+
+            # Pack the sequences into a single list
+            input_ids_list = [item["input_ids"] for item in batch]
+            position_ids_list = [item["position_ids"] for item in batch]
+            seqlens = np.array([0] + [len(ids) for ids in input_ids_list])
+
+            packed_batch = {
+                "position_ids": np.concatenate(position_ids_list, axis=0),
+                "input_ids": np.concatenate(input_ids_list, axis=0),
+                "cu_seqlens": np.cumsum(seqlens),
+                "max_seqlen": max(seqlens),
+            }
+
+            batch = super().__call__([packed_batch])
+            batch["cu_seqlens"] = batch["cu_seqlens"].to(torch.int32).squeeze()
+        else:
+            batch = super().__call__(batch)
+            batch["attention_mask"] = batch["attention_mask"].to(torch.bool)
+
+        return batch
+
+
+class NeoBERTConfig(PretrainedConfig):
+    model_type = "neobert"
+
+    # All config parameters must have a default value.
+    def __init__(
+        self,
+        hidden_size: int = 768,
+        num_hidden_layers: int = 28,
+        num_attention_heads: int = 12,
+        intermediate_size: int = 3072,
+        embedding_init_range: float = 0.02,
+        decoder_init_range: float = 0.02,
+        norm_eps: float = 1e-05,
+        vocab_size: int = 65000,
+        pad_token_id: int = 0,
+        max_length: int = 1024,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        if hidden_size % num_attention_heads != 0:
+            raise ValueError("Hidden size must be divisible by the number of heads.")
+        self.dim_head = hidden_size // num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.embedding_init_range = embedding_init_range
+        self.decoder_init_range = decoder_init_range
+        self.norm_eps = norm_eps
+        self.vocab_size = vocab_size
+        self.pad_token_id = pad_token_id
+        self.max_length = max_length
+        self.kwargs = kwargs
+
+
+class EncoderBlock(nn.Module):
+    """Transformer encoder block."""
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__()
+
+        self.config = config
+
+        # Attention
+        self.qkv = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size * 3, bias=False)
+        self.wo = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size, bias=False)
+
+        # Feedforward network
+        multiple_of = 8
+        intermediate_size = int(2 * config.intermediate_size / 3)
+        intermediate_size = multiple_of * ((intermediate_size + multiple_of - 1) // multiple_of)
+        self.ffn = SwiGLU(config.hidden_size, intermediate_size, config.hidden_size, bias=False)
+
+        # Layer norms
+        self.attention_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)
+        self.ffn_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        output_attentions: bool,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+    ):
+        # Attention
+        attn_output, attn_weights = self._att_block(
+            self.attention_norm(x), attention_mask, freqs_cis, output_attentions, max_seqlen, cu_seqlens
+        )
+
+        # Residual
+        x = x + attn_output
+
+        # Feed-forward
+        x = x + self.ffn(self.ffn_norm(x))
+
+        return x, attn_weights
+
+    def _att_block(
+        self,
+        x: torch.Tensor,
+        attention_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        output_attentions: bool,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+    ):
+        batch_size, seq_len, _ = x.shape
+
+        xq, xk, xv = self.qkv(x).view(batch_size, seq_len, self.config.num_attention_heads, self.config.dim_head * 3).chunk(3, axis=-1)
+
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+
+        # Attn block
+        attn_weights = None
+
+        # Flash attention if the tensors are packed
+        if cu_seqlens is not None:
+            attn = flash_attn_varlen_func(
+                q=xq.squeeze(0),
+                k=xk.squeeze(0),
+                v=xv.squeeze(0),
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                dropout_p=0.0,
+                causal=False,
+            )
+        # Eager attention if attention weights are needed in the output
+        elif output_attentions:
+            attn_weights = xq.permute(0, 2, 1, 3) @ xk.permute(0, 2, 3, 1) / (xq.size(-1) ** 0.5)
+            if attention_mask is not None:
+                attn_weights = attn_weights * attention_mask
+            attn_weights = attn_weights.softmax(-1)
+            attn = attn_weights @ xv.permute(0, 2, 1, 3)
+            attn = attn.transpose(1, 2)
+        # Fall back to SDPA otherwise
+        else:
+            attn = scaled_dot_product_attention(
+                query=xq.transpose(1, 2),
+                key=xk.transpose(1, 2),
+                value=xv.transpose(1, 2),
+                attn_mask=attention_mask.bool(),
+                dropout_p=0,
+            ).transpose(1, 2)
+
+        return self.wo(attn.reshape(batch_size, seq_len, self.config.num_attention_heads * self.config.dim_head)), attn_weights
+
+
+class NeoBERTPreTrainedModel(PreTrainedModel):
+    config_class = NeoBERTConfig
+    base_model_prefix = "model"
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.uniform_(-self.config.decoder_init_range, self.config.decoder_init_range)
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.uniform_(-self.config.embedding_init_range, self.config.embedding_init_range)
+
+
+class NeoBERT(NeoBERTPreTrainedModel):
+    config_class = NeoBERTConfig
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__(config)
+
+        self.config = config
+
+        self.encoder = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+
+        # Ensures freqs_cis is moved to the same devices as the model. Non-persistent buffers are not saved in the state_dict.
+        freqs_cis = precompute_freqs_cis(config.hidden_size // config.num_attention_heads, config.max_length)
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
+
+        self.transformer_encoder = nn.ModuleList()
+        for _ in range(config.num_hidden_layers):
+            self.transformer_encoder.append(EncoderBlock(config))
+
+        self.layer_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        **kwargs,
+    ):
+        # Initialize
+        hidden_states, attentions = [], []
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        # Expand and repeat: (Batch, Length) -> (Batch, Heads, Length, Length)
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).repeat(1, self.config.num_attention_heads, attention_mask.size(-1), 1)
+
+        # Checks to be done if inputs are packed sequences
+        if cu_seqlens is not None:
+            assert (
+                FLASH_ATTN_AVAILABLE
+            ), "Flash-attention is not available. Please ''pip install flash_attn'', or provide un-packed sequences."
+            assert not output_attentions, "Output attentions is not supported when sequences are packed."
+            assert max_seqlen is not None, "Missing max_seqlen. It must be provided when cu_seqlens are not None."
+            assert (input_ids if input_ids is not None else inputs_embeds).shape[
+                0
+            ] == 1, "Cumulative sequence lengths are provided but inputs are not packed."
+            assert (
+                input_ids if input_ids is not None else inputs_embeds
+            ).is_cuda, "Packing uses an implementation of flash-attention and is only supported on GPU."
+
+        # RoPE
+        freqs_cis = (
+            self.freqs_cis[position_ids]
+            if position_ids is not None
+            else self.freqs_cis[: (input_ids if input_ids is not None else inputs_embeds).shape[1]].unsqueeze(0)
+        )
+
+        # Embedding
+        x = self.encoder(input_ids) if input_ids is not None else inputs_embeds
+
+        # Transformer encoder
+        for layer in self.transformer_encoder:
+            x, attn = layer(x, attention_mask, freqs_cis, output_attentions, max_seqlen, cu_seqlens)
+            if output_hidden_states:
+                hidden_states.append(x)
+            if output_attentions:
+                attentions.append(attn)
+
+        # Final normalization layer
+        x = self.layer_norm(x)
+
+        # Return the output of the last hidden layer
+        return BaseModelOutput(
+            last_hidden_state=x,
+            hidden_states=hidden_states if output_hidden_states else None,
+            attentions=attentions if output_attentions else None,
+        )
+
+
+class NeoBERTLMHead(NeoBERTPreTrainedModel):
+    config_class = NeoBERTConfig
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__(config)
+
+        self.config = config
+
+        self.model = NeoBERT(config)
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
+
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        **kwargs,
+    ):
+
+        output = self.model.forward(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            max_seqlen=max_seqlen,
+            cu_seqlens=cu_seqlens,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+        logits = self.decoder(output.last_hidden_state)
+
+        return MaskedLMOutput(
+            hidden_states=output.hidden_states if output_hidden_states else None,
+            attentions=output.attentions if output_attentions else None,
+            logits=logits,
+        )
+
+
+class NeoBERTForSequenceClassification(NeoBERTPreTrainedModel):
+    config_class = NeoBERTConfig
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__(config)
+
+        self.config = config
+
+        self.num_labels = getattr(config, "num_labels", 2)
+        self.classifier_dropout = getattr(config, "classifier_dropout", 0.1)
+        self.classifier_init_range = getattr(config, "classifier_init_range", 0.02)
+
+        self.model = NeoBERT(config)
+
+        self.dense = nn.Linear(self.config.hidden_size, self.config.hidden_size)
+        self.dropout = nn.Dropout(self.classifier_dropout)
+        self.classifier = nn.Linear(self.config.hidden_size, self.num_labels)
+
+        self.post_init()
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.classifier_init_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        output = self.model.forward(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            max_seqlen=max_seqlen,
+            cu_seqlens=cu_seqlens,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+        hidden_states = output.last_hidden_state
+
+        x = hidden_states[:, 0, :]
+        x = self.dropout(x)
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.dropout(x)
+
+        logits = self.classifier(x)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        if not return_dict:
+            result = (logits,)
+            return ((loss,) + result) if loss is not None else result
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=output.hidden_states if output_hidden_states else None,
+            attentions=output.attentions if output_attentions else None,
+        )

model.safetensors

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

modules.json

@@ -0,0 +1,14 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  }
+]

rotary.py

@@ -0,0 +1,61 @@
+# From https://github.com/facebookresearch/llama/blob/main/llama/model.py
+
+import torch
+from typing import Tuple
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    """
+    Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+    This function calculates a frequency tensor with complex exponentials using the given dimension 'dim'
+    and the end index 'end'. The 'theta' parameter scales the frequencies.
+    The returned tensor contains complex values in complex64 data type.
+
+    Args:
+        dim (int): Dimension of the frequency tensor.
+        end (int): End index for precomputing frequencies.
+        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
+
+    Returns:
+        torch.Tensor: Precomputed frequency tensor with complex exponentials.
+    """
+
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    assert freqs_cis.shape[1:] == (x.shape[1], x.shape[-1])
+    return freqs_cis.contiguous().unsqueeze(2)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor.
+
+    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided
+    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor
+    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are
+    returned as real tensors.
+
+    Args:
+        xq (torch.Tensor): Query tensor to apply rotary embeddings.
+        xk (torch.Tensor): Key tensor to apply rotary embeddings.
+        freqs_cis (torch.Tensor): Precomputed frequency tensor for complex exponentials.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+    """
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)

sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+    "max_seq_length": 512,
+    "do_lower_case": false
+}

special_tokens_map.json

@@ -0,0 +1,40 @@
+{
+  "additional_special_tokens": [
+    "[+]"
+  ],
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.py

@@ -0,0 +1,147 @@
+from typing import List, Tuple
+from transformers import PreTrainedTokenizerFast
+import re
+from .acrps_light_stemming import stem
+from .constants import ARABIC_DIACRITICS
+
+_TATWEEL_RE = re.compile(r"\u0640")
+_ALIF_RE = re.compile(r"[آأإٱ]")
+_ALIF_MAK_RE = re.compile(r"ى")
+_TEH_MARB_RE = re.compile(r"ة")
+_ZERO_WIDTH_RE = re.compile(r"[\u200B-\u200D\u200E\u200F\uFEFF]")
+
+def separate_diacritics(text):
+    tokens = re.split(r'(\s+|\[\+\])', text)
+    processed_tokens = []
+    
+    for token in tokens:
+        if not token: 
+            continue
+        if token.isspace() or token == '[+]':
+            processed_tokens.append(token)
+            continue
+            
+        if not any(c in ARABIC_DIACRITICS for c in token):
+            processed_tokens.append(token)
+            continue
+            
+        base_chars = []
+        diac_groups = []
+        
+        for char in token:
+            if char in ARABIC_DIACRITICS:
+                if not diac_groups:
+                    base_chars.append(" ")
+                    diac_groups.append([])
+                diac_groups[-1].append(char)
+            else:
+                base_chars.append(char)
+                diac_groups.append([])
+        
+        base_word = "".join(base_chars)
+        diac_string = []
+        for group in diac_groups:
+            if group:
+                diac_string.append("".join(group))
+            else:
+                diac_string.append("◌")
+        
+        processed_tokens.append(base_word + " " + "".join(diac_string))
+    return "".join(processed_tokens)
+
+def normalize_arabic(text):
+    text = _TATWEEL_RE.sub("", text)
+    text = _ZERO_WIDTH_RE.sub("", text)
+    text = _ALIF_RE.sub("ا", text)
+    text = _ALIF_MAK_RE.sub("ي", text)
+    text = _TEH_MARB_RE.sub("ه", text)
+    return text
+
+class ArabicMorphTokenizer(PreTrainedTokenizerFast):
+    slow_tokenizer_class = None
+
+    def __init__(self, tokenizer_file=None, apply_stemming=True, **kwargs):
+        super().__init__(tokenizer_file=tokenizer_file, **kwargs)
+        self.apply_stemming = True if apply_stemming is None else bool(apply_stemming)
+
+    def _preprocess_one(self, s, do_stem):
+        if isinstance(s, (list, tuple)):
+            return [self._preprocess_one(x, do_stem) for x in s]
+        if do_stem:
+            s = stem(s, apply_diacritics=True)
+        s = normalize_arabic(s)
+        s = separate_diacritics(s)
+        return s
+
+    def _preprocess_pair(self, text, text_pair, do_stem):
+        def maybe(s):
+            return self._preprocess_one(s, do_stem) if isinstance(s, str) else s
+        if isinstance(text, (list, tuple)):
+            text = [maybe(x) for x in text]
+        else:
+            text = maybe(text)
+        if isinstance(text_pair, (list, tuple)):
+            text_pair = [maybe(x) for x in text_pair]
+        else:
+            text_pair = maybe(text_pair)
+        return text, text_pair
+
+    def _pop_flag(self, kwargs):
+        v = kwargs.pop("apply_stemming", None)
+        return self.apply_stemming if v is None else bool(v)
+
+    def __call__(self, text=None, text_pair=None, *args, **kwargs):
+        flag = self._pop_flag(kwargs)
+        if not getattr(self, "_processing", False):
+            self._processing = True
+            try:
+                text, text_pair = self._preprocess_pair(text, text_pair, flag)
+                return super().__call__(text=text, text_pair=text_pair, *args, **kwargs)
+            finally:
+                self._processing = False
+        return super().__call__(text=text, text_pair=text_pair, *args, **kwargs)
+
+    def encode(self, text, text_pair=None, *args, **kwargs):
+        flag = self._pop_flag(kwargs)
+        if not getattr(self, "_processing", False):
+            self._processing = True
+            try:
+                text, text_pair = self._preprocess_pair(text, text_pair, flag)
+                return super().encode(text, text_pair, *args, **kwargs)
+            finally:
+                self._processing = False
+        return super().encode(text, text_pair, *args, **kwargs)
+
+    def encode_plus(self, text=None, text_pair=None, *args, **kwargs):
+        flag = self._pop_flag(kwargs)
+        if not getattr(self, "_processing", False):
+            self._processing = True
+            try:
+                text, text_pair = self._preprocess_pair(text, text_pair, flag)
+                return super().encode_plus(text=text, text_pair=text_pair, *args, **kwargs)
+            finally:
+                self._processing = False
+        return super().encode_plus(text=text, text_pair=text_pair, *args, **kwargs)
+
+    def batch_encode_plus(self, batch_text_or_text_pairs=None, *args, **kwargs):
+        flag = self._pop_flag(kwargs)
+        if not getattr(self, "_processing", False):
+            self._processing = True
+            try:
+                data = batch_text_or_text_pairs
+                if isinstance(data, (list, tuple)):
+                    new_data = []
+                    for item in data:
+                        if isinstance(item, (list, tuple)) and len(item) == 2:
+                            new_data.append(self._preprocess_pair(item[0], item[1], flag))
+                        else:
+                            new_data.append(self._preprocess_one(item, flag))
+                    batch_text_or_text_pairs = new_data
+                return super().batch_encode_plus(batch_text_or_text_pairs=batch_text_or_text_pairs, *args, **kwargs)
+            finally:
+                self._processing = False
+        return super().batch_encode_plus(batch_text_or_text_pairs=batch_text_or_text_pairs, *args, **kwargs)
+
+    def preprocess(self, text, apply_stemming=True):
+        flag = self.apply_stemming if apply_stemming is None else bool(apply_stemming)
+        return self._preprocess_one(text, flag)

tokenizer_config.json

@@ -0,0 +1,74 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "[+]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "[+]"
+  ],
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenizer.ArabicMorphTokenizer",
+      null
+    ]
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "do_lower_case": false,
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "ArabicMorphTokenizer",
+  "trust_remote_code": true,
+  "unk_token": "[UNK]"
+}