Upload model.py

model.py

@@ -0,0 +1,404 @@
+"""
+================================================
+Arabic Diacritization - mishkala
+نموذج التشكيل العربي التلقائي
+https://huggingface.co/flokymind/mishkala
+================================================
+المتطلبات:
+    pip install torch pytorch-crf huggingface_hub
+================================================
+"""
+
+# ── المتطلبات ─────────────────────────────────
+import subprocess, sys
+
+def install(pkg):
+    subprocess.check_call([sys.executable, "-m", "pip", "install", pkg, "-q"])
+
+try:
+    import torchcrf
+except ImportError:
+    install("pytorch-crf")
+
+try:
+    from huggingface_hub import hf_hub_download
+except ImportError:
+    install("huggingface_hub")
+
+# ── الاستيرادات ───────────────────────────────
+import json, math, re
+from pathlib import Path
+from typing import Dict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchcrf import CRF
+from huggingface_hub import hf_hub_download
+
+
+# ════════════════════════════════════════════
+# 1. الثوابت
+# ════════════════════════════════════════════
+
+REPO_ID = "flokymind/mishkala"
+
+DIACRITICS_SET = {
+    '\u064e', '\u064b', '\u064f', '\u064c',
+    '\u0650', '\u064d', '\u0651', '\u0652',
+}
+
+SPECIAL_TOKENS = {'PAD': 0, 'UNK': 1, 'BOS': 2, 'EOS': 3, 'MASK': 4, ' ': 5}
+
+DIACRITIC_CLASSES = [
+    'NO_DIACRITIC', 'FATHA', 'FATHATAN', 'DAMMA', 'DAMMATAN',
+    'KASRA', 'KASRATAN', 'SUKUN', 'SHADDA',
+    'SHADDA_FATHA', 'SHADDA_FATHATAN', 'SHADDA_DAMMA',
+    'SHADDA_DAMMATAN', 'SHADDA_KASRA', 'SHADDA_KASRATAN',
+]
+
+DIACRITIC_MAP = {
+    'NO_DIACRITIC':      '',
+    'FATHA':             '\u064e',
+    'FATHATAN':          '\u064b',
+    'DAMMA':             '\u064f',
+    'DAMMATAN':          '\u064c',
+    'KASRA':             '\u0650',
+    'KASRATAN':          '\u064d',
+    'SUKUN':             '\u0652',
+    'SHADDA':            '\u0651',
+    'SHADDA_FATHA':      '\u0651\u064e',
+    'SHADDA_FATHATAN':   '\u0651\u064b',
+    'SHADDA_DAMMA':      '\u0651\u064f',
+    'SHADDA_DAMMATAN':   '\u0651\u064c',
+    'SHADDA_KASRA':      '\u0651\u0650',
+    'SHADDA_KASRATAN':   '\u0651\u064d',
+}
+
+
+# ════════════════════════════════════════════
+# 2. التوكنايزر
+# ════════════════════════════════════════════
+
+class ArabicTokenizer:
+    def __init__(self):
+        self.char_to_id: Dict[str, int] = {}
+        self.id_to_char: Dict[int, str] = {}
+        self.vocab_size: int = 0
+
+    def encode(self, text, max_length=512, padding=True):
+        ids = [SPECIAL_TOKENS['BOS']]
+        for ch in text:
+            if ch in DIACRITICS_SET:
+                continue
+            ids.append(self.char_to_id.get(ch, SPECIAL_TOKENS['UNK']))
+        ids.append(SPECIAL_TOKENS['EOS'])
+
+        attention_mask = [1] * len(ids)
+
+        if len(ids) > max_length:
+            ids            = ids[:max_length]
+            attention_mask = attention_mask[:max_length]
+        elif padding:
+            pad_len         = max_length - len(ids)
+            ids            += [SPECIAL_TOKENS['PAD']] * pad_len
+            attention_mask += [0] * pad_len
+
+        return ids, attention_mask
+
+    @classmethod
+    def load(cls, path):
+        data           = json.loads(Path(path).read_text(encoding='utf-8'))
+        tok            = cls()
+        tok.char_to_id = data['char_to_id']
+        tok.id_to_char = {int(v): k for k, v in data['char_to_id'].items()}
+        tok.vocab_size = data['vocab_size']
+        print(f"✅ التوكنايزر: {tok.vocab_size} رمز")
+        return tok
+
+
+# ════════════════════════════════════════════
+# 3. مكونات النموذج
+# ════════════════════════════════════════════
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps   = eps
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
+        return self.scale * x / rms
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len=4096):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+        t     = torch.arange(max_seq_len).float()
+        freqs = torch.outer(t, inv_freq)
+        emb   = torch.cat([freqs, freqs], dim=-1)
+        self.register_buffer('cos_cached', emb.cos())
+        self.register_buffer('sin_cached', emb.sin())
+
+    def forward(self, x, seq_len):
+        return (
+            self.cos_cached[:seq_len].unsqueeze(0),
+            self.sin_cached[:seq_len].unsqueeze(0),
+        )
+
+
+def rotate_half(x):
+    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rope(q, k, cos, sin):
+    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim, expansion=4):
+        super().__init__()
+        hidden         = int(dim * expansion * 2 / 3)
+        hidden         = (hidden + 7) // 8 * 8
+        self.gate_proj = nn.Linear(dim, hidden, bias=False)
+        self.up_proj   = nn.Linear(dim, hidden, bias=False)
+        self.down_proj = nn.Linear(hidden, dim, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class MambaBlock(nn.Module):
+    def __init__(self, dim, d_state=16, d_conv=4, expand=2):
+        super().__init__()
+        self.d_inner  = int(dim * expand)
+        self.in_proj  = nn.Linear(dim, self.d_inner * 2, bias=False)
+        self.conv1d   = nn.Conv1d(self.d_inner, self.d_inner, d_conv,
+                                  padding=d_conv-1, groups=self.d_inner, bias=True)
+        self.out_proj = nn.Linear(self.d_inner, dim, bias=False)
+        self.x_proj   = nn.Linear(self.d_inner, d_state * 2 + 1, bias=False)
+        self.dt_proj  = nn.Linear(1, self.d_inner, bias=True)
+        A             = torch.arange(1, d_state+1).float().unsqueeze(0).expand(self.d_inner, -1)
+        self.A_log    = nn.Parameter(torch.log(A))
+        self.D        = nn.Parameter(torch.ones(self.d_inner))
+        self.norm     = RMSNorm(dim)
+
+    def ssm(self, x):
+        dt = F.softplus(self.dt_proj(self.x_proj(x)[..., :1]))
+        return x * self.D + torch.cumsum(x * dt, dim=1) * 0.1
+
+    def forward(self, x):
+        residual = x
+        x        = self.norm(x)
+        xz       = self.in_proj(x)
+        x_ssm, z = xz.chunk(2, dim=-1)
+        x_conv   = self.conv1d(x_ssm.transpose(1,2))[..., :x_ssm.shape[1]].transpose(1,2)
+        y        = self.ssm(F.silu(x_conv)) * F.silu(z)
+        return self.out_proj(y) + residual
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, n_heads, max_len=4096, dropout=0.1):
+        super().__init__()
+        self.n_heads  = n_heads
+        self.head_dim = dim // n_heads
+        self.q_proj   = nn.Linear(dim, dim, bias=False)
+        self.k_proj   = nn.Linear(dim, dim, bias=False)
+        self.v_proj   = nn.Linear(dim, dim, bias=False)
+        self.o_proj   = nn.Linear(dim, dim, bias=False)
+        self.rope     = RotaryEmbedding(self.head_dim, max_len)
+        self.ffn      = SwiGLU(dim)
+        self.norm1    = RMSNorm(dim)
+        self.norm2    = RMSNorm(dim)
+        self.dropout  = nn.Dropout(dropout)
+
+    def attention(self, x, mask=None):
+        B, L, D = x.shape
+        q = self.q_proj(x).view(B,L,self.n_heads,self.head_dim).transpose(1,2)
+        k = self.k_proj(x).view(B,L,self.n_heads,self.head_dim).transpose(1,2)
+        v = self.v_proj(x).view(B,L,self.n_heads,self.head_dim).transpose(1,2)
+        cos, sin = self.rope(x, L)
+        cos = cos.unsqueeze(1).expand_as(q)
+        sin = sin.unsqueeze(1).expand_as(q)
+        q, k = apply_rope(q, k, cos, sin)
+        scores = torch.matmul(q, k.transpose(-2,-1)) / math.sqrt(self.head_dim)
+        if mask is not None:
+            scores = scores.masked_fill(
+                ~mask.unsqueeze(1).unsqueeze(2).bool(), float('-inf')
+            )
+        attn = self.dropout(F.softmax(scores, dim=-1))
+        out  = torch.matmul(attn, v).transpose(1,2).contiguous().view(B,L,D)
+        return self.o_proj(out)
+
+    def forward(self, x, mask=None):
+        x = x + self.dropout(self.attention(self.norm1(x), mask))
+        x = x + self.dropout(self.ffn(self.norm2(x)))
+        return x
+
+
+class ArabicDiacritizerModel(nn.Module):
+    def __init__(self, vocab_size=50, dim=320, mamba_layers=4,
+                 transformer_layers=8, n_heads=8, num_labels=15,
+                 max_seq_len=4096, dropout=0.15, d_state=16):
+        super().__init__()
+        self.num_labels         = num_labels
+        self.embedding          = nn.Embedding(vocab_size, dim, padding_idx=0)
+        self.emb_norm           = RMSNorm(dim)
+        self.dropout            = nn.Dropout(dropout)
+        self.mamba_layers       = nn.ModuleList([
+            MambaBlock(dim, d_state) for _ in range(mamba_layers)
+        ])
+        self.transformer_layers = nn.ModuleList([
+            TransformerBlock(dim, n_heads, max_seq_len, dropout)
+            for _ in range(transformer_layers)
+        ])
+        self.final_norm         = RMSNorm(dim)
+        self.classifier         = nn.Linear(dim, num_labels)
+        self.crf                = CRF(num_labels, batch_first=True)
+
+    def forward(self, input_ids, attention_mask=None, labels=None):
+        x = self.dropout(self.emb_norm(self.embedding(input_ids)))
+        for m in self.mamba_layers:
+            x = m(x)
+        for t in self.transformer_layers:
+            x = t(x, attention_mask)
+        emissions = self.classifier(self.final_norm(x))
+        mask      = (attention_mask.bool() if attention_mask is not None
+                     else torch.ones(emissions.shape[:2],
+                                     dtype=torch.bool, device=emissions.device))
+        return {
+            'predictions': self.crf.decode(emissions, mask=mask),
+            'emissions':   emissions,
+        }
+
+
+# ════════════════════════════════════════════
+# 4. تحميل النموذج من HuggingFace
+# ════════════════════════════════════════════
+
+def load_mishkala(repo_id: str = REPO_ID, device: str = None):
+    """
+    تحميل نموذج مِشكالة من HuggingFace
+
+    مثال:
+        model, tokenizer, device = load_mishkala()
+    """
+    if device is None:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    device = torch.device(device)
+
+    print(f"📥 تحميل مِشكالة من {repo_id}...")
+
+    tokenizer_path = hf_hub_download(repo_id=repo_id, filename="tokenizer.json")
+    tokenizer      = ArabicTokenizer.load(tokenizer_path)
+
+    ckpt_path    = hf_hub_download(repo_id=repo_id, filename="mishkala.pt")
+    ckpt         = torch.load(ckpt_path, map_location=device)
+    model_config = ckpt['config']
+    model        = ArabicDiacritizerModel(**model_config).to(device)
+    model.load_state_dict(ckpt['model_state_dict'])
+    model.eval()
+
+    params = sum(p.numel() for p in model.parameters())
+    print(f"✅ النموذج جاهز | Step: {ckpt['step']:,} | DER: {ckpt['der']*100:.2f}%")
+    print(f"   {device} | {params:,} معلمة")
+
+    return model, tokenizer, device
+
+
+# ════════════════════════════════════════════
+# 5. دالة التشكيل
+# ════════════════════════════════════════════
+
+def tashkeel(
+    text:      str,
+    model:     ArabicDiacritizerModel = None,
+    tokenizer: ArabicTokenizer        = None,
+    device:    torch.device           = None,
+    max_chunk: int                    = 400,
+) -> str:
+    """
+    شكّل أي نص عربي تلقائياً
+
+    المعاملات:
+        text      : النص العربي المراد تشكيله
+        model     : النموذج (يُحمَّل تلقائياً إذا لم يُعطَ)
+        tokenizer : التوكنايزر (يُحمَّل تلقائياً إذا لم يُعطَ)
+        device    : الجهاز cuda/cpu
+        max_chunk : الحد الأقصى لطول القطعة الواحدة
+
+    المخرج:
+        النص مشكّلاً كاملاً
+
+    مثال:
+        model, tokenizer, device = load_mishkala()
+        result = tashkeel("كان الفيلسوف يرى أن العقل مرآة", model, tokenizer, device)
+        print(result)
+        # كَانَ الْفَيْلَسُوفُ يَرَى أَنَّ الْعَقْلَ مِرْآةٌ
+    """
+    # تحميل تلقائي إذا لم يُعطَ نموذج
+    global _default_model, _default_tokenizer, _default_device
+    if model is None:
+        if '_default_model' not in globals():
+            _default_model, _default_tokenizer, _default_device = load_mishkala()
+        model, tokenizer, device = _default_model, _default_tokenizer, _default_device
+
+    # إزالة التشكيل الموجود
+    clean = ''.join(c for c in text if c not in DIACRITICS_SET)
+
+    # تقسيم النص على الجمل
+    sentences = re.split(r'([.،؟!\n])', clean)
+    chunks, current = [], ""
+    for part in sentences:
+        if len(current) + len(part) > max_chunk and current:
+            chunks.append(current.strip())
+            current = part
+        else:
+            current += part
+    if current.strip():
+        chunks.append(current.strip())
+
+    results = []
+    for chunk in chunks:
+        if not chunk.strip():
+            results.append(chunk)
+            continue
+
+        input_ids, attention_mask = tokenizer.encode(chunk, max_length=512, padding=True)
+        ids_t  = torch.tensor([input_ids],      dtype=torch.long).to(device)
+        mask_t = torch.tensor([attention_mask], dtype=torch.long).to(device)
+
+        with torch.no_grad():
+            out = model(ids_t, mask_t)
+
+        pred_labels  = out['predictions'][0]
+        chars        = [c for c in chunk if c not in DIACRITICS_SET]
+        result_chars = []
+
+        for i, char in enumerate(chars):
+            result_chars.append(char)
+            label_idx = i + 1
+            if label_idx < len(pred_labels):
+                diacritic = DIACRITIC_MAP.get(
+                    DIACRITIC_CLASSES[pred_labels[label_idx]], ''
+                )
+                result_chars.append(diacritic)
+
+        results.append(''.join(result_chars))
+
+    return ''.join(results)
+
+
+# ════════════════════════════════════════════
+# 6. التشغيل المباشر
+# ════════════════════════════════════════════
+
+if __name__ == "__main__":
+    model, tokenizer, device = load_mishkala()
+
+    text = "الإنسان بين العقل والغريزة"
+    print(f"\n✨ {tashkeel(text, model, tokenizer, device)}")