Rename start.sh to app.py

app.py

@@ -0,0 +1,231 @@
+import os
+import time
+import math
+import jax
+import jax.numpy as jnp
+import flax.linen as nn
+import orbax.checkpoint as ocp
+from transformers import AutoTokenizer
+from huggingface_hub import snapshot_download
+
+# ════════════════════════════════════════════════════════════════════════
+# 1. إعدادات النموذج والمتغيرات
+# ════════════════════════════════════════════════════════════════════════
+REPO_ID = "AutoCode-Ai/HTDN-Model"
+TOKENIZER_NAME = "bert-base-multilingual-cased"
+
+D_MODEL = 1024
+D_U = 512
+D_V = 512
+HIDDEN_REACT = 3072
+N_ATTRACTORS = 22_000
+NUM_HEADS = 8
+NUM_LAYERS = 4
+MAX_LEN = 64
+COMPUTE_DTYPE = jnp.float16
+
+# ════════════════════════════════════════════════════════════════════════
+# 2. معمارية HTDN الدقيقة المطابقة للأوزان
+# ════════════════════════════════════════════════════════════════════════
+def rope_frequencies(d_head: int, base: float = 10000.0) -> jnp.ndarray:
+    inv_freq = 1.0 / (base ** (jnp.arange(0, d_head, 2, dtype=jnp.float32) / d_head))
+    return inv_freq
+
+def apply_rope(x: jnp.ndarray, positions: jnp.ndarray, inv_freq: jnp.ndarray) -> jnp.ndarray:
+    freqs = jnp.einsum("l,d->ld", positions.astype(jnp.float32), inv_freq)
+    cos = jnp.cos(freqs)[None, None, :, :]
+    sin = jnp.sin(freqs)[None, None, :, :]
+    x1, x2 = jnp.split(x.astype(jnp.float32), 2, axis=-1)
+    return jnp.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1).astype(x.dtype)
+
+def _linear_recurrence_combine(l, r):
+    a_l, b_l = l
+    a_r, b_r = r
+    return a_l * a_r, a_r * b_l + b_r
+
+class DynamicsLayer(nn.Module):
+    d_u: int
+    d_v: int
+    n_attractors: int
+    num_heads: int
+
+    @nn.compact
+    def __call__(self, x: jnp.ndarray, steps: int) -> jnp.ndarray:
+        u = nn.Dense(self.d_u, dtype=COMPUTE_DTYPE)(x)
+        v = nn.Dense(self.d_v, dtype=COMPUTE_DTYPE)(x)
+
+        log_Du = self.param("log_Du", nn.initializers.normal(stddev=0.01), (self.num_heads,))
+        log_Dv = self.param("log_Dv", nn.initializers.normal(stddev=0.01), (self.num_heads,))
+        Du = nn.softplus(log_Du).mean()
+        Dv = nn.softplus(log_Dv).mean()
+
+        attr_u = nn.Dense(self.n_attractors, dtype=COMPUTE_DTYPE)(u)
+        attr_v = nn.Dense(self.n_attractors, dtype=COMPUTE_DTYPE)(v)
+        back_u = nn.Dense(self.d_u, dtype=COMPUTE_DTYPE)(attr_u)
+        back_v = nn.Dense(self.d_v, dtype=COMPUTE_DTYPE)(attr_v)
+
+        SAFE_BOUND = 30.0
+        u0 = jnp.clip(u.astype(jnp.float32), -SAFE_BOUND, SAFE_BOUND)
+        v0 = jnp.clip(v.astype(jnp.float32), -SAFE_BOUND, SAFE_BOUND)
+        back_u = jnp.clip(back_u.astype(jnp.float32), -SAFE_BOUND, SAFE_BOUND)
+        back_v = jnp.clip(back_v.astype(jnp.float32), -SAFE_BOUND, SAFE_BOUND)
+
+        dt = 0.1 / max(1, steps)
+        a_u = jnp.clip(1.0 - dt * Du.astype(jnp.float32), -0.99, 1.0)
+        a_v = jnp.clip(1.0 - dt * Dv.astype(jnp.float32), -0.99, 1.0)
+        b_u = dt * back_u
+        b_v = dt * back_v
+
+        A_u_seq = jnp.broadcast_to(a_u, (steps,) + u0.shape)
+        B_u_seq = jnp.broadcast_to(b_u, (steps,) + u0.shape)
+        A_v_seq = jnp.broadcast_to(a_v, (steps,) + v0.shape)
+        B_v_seq = jnp.broadcast_to(b_v, (steps,) + v0.shape)
+
+        cumA_u, cumB_u = jax.lax.associative_scan(_linear_recurrence_combine, (A_u_seq, B_u_seq), axis=0)
+        cumA_v, cumB_v = jax.lax.associative_scan(_linear_recurrence_combine, (A_v_seq, B_v_seq), axis=0)
+
+        u_final = jnp.clip(cumA_u[-1] * u0 + cumB_u[-1], -SAFE_BOUND, SAFE_BOUND)
+        v_final = jnp.clip(cumA_v[-1] * v0 + cumB_v[-1], -SAFE_BOUND, SAFE_BOUND)
+
+        return jnp.concatenate([u_final, v_final], axis=-1).astype(COMPUTE_DTYPE)
+
+class CausalSelfAttention(nn.Module):
+    num_heads: int
+    d_model: int
+
+    @nn.compact
+    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
+        B, L, D = x.shape
+        d_head = self.d_model // self.num_heads
+        Q = nn.Dense(self.d_model, dtype=COMPUTE_DTYPE)(x).reshape(B, L, self.num_heads, d_head).transpose(0, 2, 1, 3)
+        K = nn.Dense(self.d_model, dtype=COMPUTE_DTYPE)(x).reshape(B, L, self.num_heads, d_head).transpose(0, 2, 1, 3)
+        V = nn.Dense(self.d_model, dtype=COMPUTE_DTYPE)(x).reshape(B, L, self.num_heads, d_head).transpose(0, 2, 1, 3)
+
+        inv_freq = rope_frequencies(d_head)
+        positions = jnp.arange(L)
+        Q = apply_rope(Q, positions, inv_freq)
+        K = apply_rope(K, positions, inv_freq)
+
+        scores = jnp.matmul(Q, K.transpose(0, 1, 3, 2)) / math.sqrt(d_head)
+        idx = jnp.arange(L)
+        causal_mask = jnp.where(idx[:, None] >= idx[None, :], 0.0, -1e9)
+        scores = scores + causal_mask[None, None, :, :]
+        attn = nn.softmax(scores, axis=-1)
+        out = jnp.matmul(attn, V).transpose(0, 2, 1, 3).reshape(B, L, self.d_model)
+
+        return nn.Dense(self.d_model, dtype=COMPUTE_DTYPE)(out)
+
+class SwiGLUFFN(nn.Module):
+    d_model: int
+    hidden_dim: int
+
+    @nn.compact
+    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
+        swiglu_hidden = max(1, int(self.hidden_dim * 2 / 3))
+        gate = nn.Dense(swiglu_hidden, dtype=COMPUTE_DTYPE)(x)
+        up = nn.Dense(swiglu_hidden, dtype=COMPUTE_DTYPE)(x)
+        return nn.Dense(self.d_model, dtype=COMPUTE_DTYPE)(nn.swish(gate) * up)
+
+class HTDNBlock(nn.Module):
+    d_model: int
+    d_u: int
+    d_v: int
+    hidden_react: int
+    n_attractors: int
+    num_heads: int
+
+    @nn.compact
+    def __call__(self, x: jnp.ndarray, steps: int) -> jnp.ndarray:
+        x = x + DynamicsLayer(self.d_u, self.d_v, self.n_attractors, self.num_heads)(nn.RMSNorm(dtype=COMPUTE_DTYPE)(x), steps)
+        x = x + CausalSelfAttention(self.num_heads, self.d_model)(nn.RMSNorm(dtype=COMPUTE_DTYPE)(x))
+        x = x + SwiGLUFFN(self.d_model, self.hidden_react)(nn.RMSNorm(dtype=COMPUTE_DTYPE)(x))
+        return x
+
+class HTDNModel(nn.Module):
+    vocab_size: int
+    d_model: int
+    d_u: int
+    d_v: int
+    hidden_react: int
+    n_attractors: int
+    num_heads: int
+    num_layers: int
+    max_len: int
+
+    @nn.compact
+    def __call__(self, input_ids: jnp.ndarray, steps: int) -> jnp.ndarray:
+        embed = nn.Embed(self.vocab_size, self.d_model, dtype=COMPUTE_DTYPE)
+        x = embed(input_ids)
+        for _ in range(self.num_layers):
+            x = HTDNBlock(self.d_model, self.d_u, self.d_v, self.hidden_react, self.n_attractors, self.num_heads)(x, steps)
+        x = nn.RMSNorm(dtype=COMPUTE_DTYPE)(x)
+        logits = jnp.matmul(x.astype(jnp.float32), embed.embedding.astype(jnp.float32).T)
+        return logits
+
+# ════════════════════════════════════════════════════════════════════════
+# 3. التحميل الآمن والأستدلال
+# ════════════════════════════════════════════════════════════════════════
+print("⏳ جاري تنزيل أوزان المستودع من Hugging Face Hub...")
+repo_path = snapshot_download(repo_id=REPO_ID)
+
+print("⏳ جاري إعداد التوكنايزر والهيكل البرمجي...")
+tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
+PAD_ID = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+
+model = HTDNModel(
+    vocab_size=tokenizer.vocab_size,
+    d_model=D_MODEL,
+    d_u=D_U,
+    d_v=D_V,
+    hidden_react=HIDDEN_REACT,
+    n_attractors=N_ATTRACTORS,
+    num_heads=NUM_HEADS,
+    num_layers=NUM_LAYERS,
+    max_len=MAX_LEN
+)
+
+rng_key = jax.random.PRNGKey(0)
+dummy_input = jnp.zeros((1, MAX_LEN), dtype=jnp.int32)
+params_template = model.init(rng_key, dummy_input, steps=24)["params"]
+
+print("⏳ جاري استعادة أوزان Orbax من المسار المحلي...")
+checkpointer = ocp.StandardCheckpointer()
+restored_params = checkpointer.restore(repo_path, target=params_template)
+print("✅ تم استعادة أوزان النموذج بنجاح!")
+
+def generate_text(prompt: str, max_new_tokens: int = 25):
+    ids = tokenizer(prompt, add_special_tokens=True, return_tensors="np")["input_ids"][0].tolist()
+    ids = ids[:MAX_LEN]
+
+    print(f"\n📝 النص المدخل: {prompt}")
+    print("🤖 التوليد: ", end="", flush=True)
+
+    for _ in range(max_new_tokens):
+        if len(ids) >= MAX_LEN:
+            break
+        cur_len = len(ids)
+        padded = ids + [PAD_ID] * (MAX_LEN - cur_len)
+        input_arr = jnp.array(padded, dtype=jnp.int32)[None, :]
+
+        logits = model.apply({"params": restored_params}, input_arr, steps=24)
+        next_id = int(jnp.argmax(logits[0, cur_len - 1]))
+
+        if next_id == PAD_ID:
+            break
+        ids.append(next_id)
+
+        token_str = tokenizer.decode([next_id], skip_special_tokens=True)
+        if token_str.startswith("##"):
+            print(token_str[2:], end="", flush=True)
+        else:
+            print(" " + token_str, end="", flush=True)
+        time.sleep(0.03)
+    print("\n" + "─" * 50)
+
+if __name__ == "__main__":
+    prompts = [
+        "تعتبر معمارية الذكاء الاصطناعي",
+        "تطوير النماذج باستخدام بايثون"
+    ]
+    for p in prompts:
+        generate_text(p)

start.sh

@@ -1,20 +0,0 @@
-#!/usr/bin/env bash
-set -e
-
-echo "🚀 Starting llama.cpp server on port 7860..."
-/app/llama-server \
-  -m /app/gemma-4-31B-it-Q4_0.gguf \
-  --mmproj /app/mmproj-BF16.gguf \
-  --host 0.0.0.0 \
-  --port 7860 \
-  -t 2 \
-  --cache-type-k q8_0 \
-  --cache-type-v iq4_nl \
-  -c 56800 \
-  -n 38912 &
-
-echo "🌐 Starting Chat UI on port 3000..."
-node /app/server.js &
-
-# Wait for both
-wait