Upload folder using huggingface_hub

README.md

@@ -1,3 +1,187 @@
----
-license: apache-2.0
----
+---
+language:
+  - fr
+license: apache-2.0
+tags:
+  - text-generation
+  - french
+  - causal-lm
+  - sft
+  - chatml
+  - from-scratch
+  - transformer
+  - recurrent-depth
+  - rope
+  - swiglu
+  - rmsnorm
+  - gqa
+library_name: transformers
+pipeline_tag: text-generation
+model-index:
+  - name: Klovis-442M-french-130426
+    results: []
+---
+
+# Klovis-442M — French Language Model
+
+A **443M-parameter** French language model using a **Recurrent-Depth Transformer** architecture, fully designed and trained from scratch by [Eric Houzelle](https://huggingface.co/Houzeric).
+
+Every component — architecture, training pipeline, and inference engine — was written in PyTorch without relying on any pre-trained weights or third-party model code. Trained end-to-end on **a single NVIDIA L40S GPU**.
+
+---
+
+## Key Facts
+
+| | |
+|:--|:--|
+| **Parameters** | ~443M |
+| **Architecture** | Recurrent-Depth Transformer (RDT) |
+| **Language** | French |
+| **Tokenizer** | CamemBERT (`camembert-base`, 32k vocab) |
+| **Context window** | 512 tokens |
+| **Chat format** | ChatML |
+| **Training hardware** | 1x NVIDIA L40S |
+| **SFT val_loss** | 1.561431465409791 |
+| **License** | Apache 2.0 |
+
+---
+
+## Quick Start
+
+```bash
+pip install transformers torch safetensors sentencepiece
+```
+
+### Conversational Mode (ChatML)
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "Klovis-ai/Klovis-442M-french-130426"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
+
+prompt = (
+    "<|system|>\n"
+    "Tu es un assistant intelligent. Réponds directement et précisément.<|end|>\n"
+    "<|user|>\n"
+    "Que peut-on visiter à Paris ?<|end|>\n"
+    "<|assistant|>\n"
+)
+
+inputs = tokenizer(prompt, return_tensors="pt")
+outputs = model.generate(
+    inputs["input_ids"],
+    max_new_tokens=200,
+    temperature=0.5,
+    top_p=0.85,
+    top_k=50,
+    do_sample=True,
+)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+### Special Tokens
+
+| Token | Role |
+|:------|:-----|
+| `<\|system\|>` | Start of system message |
+| `<\|user\|>` | Start of user message |
+| `<\|assistant\|>` | Start of assistant response |
+| `<\|end\|>` | End of turn |
+
+---
+
+## Architecture — Recurrent-Depth Transformer
+
+Klovis uses a **Recurrent-Depth Transformer** (inspired by OpenMythos/Parcae, Prairie et al. 2026), where a single Transformer block is applied iteratively:
+
+```
+Input → [Prelude: 4 layers] → [Shared Block x 8 steps] → [Coda: 4 layers] → Output
+```
+
+| Component | Detail |
+|:----------|:-------|
+| Embedding dim | 1536 |
+| Prelude layers | 4 (unique) |
+| Coda layers | 4 (unique) |
+| Recurrent steps | 8 |
+| Query heads | 24 |
+| KV heads (GQA) | 6 |
+| FFN hidden dim | 8192 |
+| FFN activation | **SwiGLU** |
+| Normalization | **RMSNorm** (pre-norm) |
+| Position encoding | **RoPE** |
+| LTI Injection | Linear Time-Invariant state coupling |
+| ACT Halting | Adaptive Computation Time (per-position early exit) |
+| Depth LoRA | Rank-32 adapters per recurrent step |
+| Weight tying | Input embeddings <-> output projection |
+
+### Why Recurrent-Depth?
+
+The RDT architecture achieves **442M effective parameters** with significantly fewer unique parameters by reusing a single Transformer block across 8 depth steps. This enables:
+- **Deeper reasoning** through iterative refinement
+- **Dynamic compute** via ACT (positions that are "easy" halt early)
+- **Step specialization** via Depth LoRA adapters
+
+---
+
+## Training
+
+### Phase 1 — Pre-training
+
+| | |
+|:--|:--|
+| **Data** | French texts (Wikipedia FR, FineWeb-2, CulturaX, C4) |
+| **Optimizer** | AdamW (beta1=0.9, beta2=0.95, wd=0.1) |
+| **Scheduler** | Linear warmup -> cosine decay |
+| **Precision** | Mixed precision (bfloat16) |
+| **Block size** | 512 tokens |
+
+### Phase 2 — Supervised Fine-Tuning (SFT)
+
+Fine-tuned on curated French conversational datasets across **25 epochs**, with prompt masking (only assistant tokens contribute to the loss).
+
+| | |
+|:--|:--|
+| **Epochs** | 25 |
+| **Learning rate** | 2e-5 |
+| **Effective batch size** | 128 (16 x 8 gradient accumulation) |
+| **Final val_loss** | 1.561431465409791 |
+| **Chat format** | ChatML |
+
+#### SFT Datasets
+
+| Dataset | Samples |
+|:--------|:--------|
+| [kurakurai/luth-sft](https://huggingface.co/datasets/kurakurai/luth-sft) (multiple subsets) | ~300k |
+| [angeluriot/french_instruct](https://huggingface.co/datasets/angeluriot/french_instruct) | 20k |
+| [Houzeric/everyday-conversations](https://huggingface.co/datasets/Houzeric/everyday-conversations) | Custom |
+| [Houzeric/physics-FR-qa-dataset](https://huggingface.co/datasets/Houzeric/physics-FR-qa-dataset) | 5k |
+| [le-leadboard/MATH_LVL5_fr](https://huggingface.co/datasets/le-leadboard/MATH_LVL5_fr) | 1.4k |
+
+---
+
+## What to Expect
+
+With ~443M parameters, Klovis produces:
+- Grammatically correct, fluent French
+- Well-structured responses (numbered lists, bold headers, paragraphs)
+- Relevant vocabulary across science, culture, history
+- Coherent conversational interactions
+
+Known limitations:
+- Factual responses may be **incorrect or fabricated** (hallucinations)
+- Mathematical reasoning is limited
+- Context limited to 512 tokens
+- French only
+
+> **This model demonstrates what a single developer can achieve with a novel architecture at small scale.** It is not intended to replace larger models for production use.
+
+---
+
+## License
+
+Apache 2.0 — [https://www.apache.org/licenses/LICENSE-2.0](https://www.apache.org/licenses/LICENSE-2.0)
+
+**Designed and trained by Eric Houzelle.**

config.json

@@ -0,0 +1,32 @@
+{
+  "model_type": "minigpt",
+  "architectures": [
+    "MiniGPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_minigpt.MiniGPTConfig",
+    "AutoModel": "modeling_minigpt.MiniGPTForCausalLM",
+    "AutoModelForCausalLM": "modeling_minigpt.MiniGPTForCausalLM"
+  },
+  "vocab_size": 32010,
+  "block_size": 512,
+  "embed_dim": 1536,
+  "depth": 14,
+  "heads": 24,
+  "num_kv_heads": 6,
+  "dropout": 0.0,
+  "hidden_dim": 8192,
+  "weight_sharing": "recurrent_depth",
+  "use_rope": true,
+  "use_gradient_checkpointing": false,
+  "num_prelude_layers": 4,
+  "num_coda_layers": 4,
+  "num_recurrent_steps": 8,
+  "use_lti_injection": true,
+  "use_act_halting": true,
+  "act_halt_threshold": 0.99,
+  "depth_lora_rank": 32,
+  "torch_dtype": "float32",
+  "tie_word_embeddings": true,
+  "tokenizer_class": "CamembertTokenizer"
+}

configuration_minigpt.py

@@ -0,0 +1,70 @@
+from transformers import PretrainedConfig
+
+
+class MiniGPTConfig(PretrainedConfig):
+    """
+    Configuration pour le modèle MiniGPT.
+    
+    Cette classe hérite de PretrainedConfig pour être compatible avec
+    l'écosystème Hugging Face.
+
+    Recurrent-Depth Transformer (RDT) — activé via weight_sharing="recurrent_depth":
+        Architecture Prelude → Recurrent Block (boucle T fois) → Coda.
+        Inspiré de OpenMythos / Parcae (Prairie et al., 2026).
+
+        Paramètres RDT:
+            num_prelude_layers: couches uniques avant la boucle (encodage initial)
+            num_coda_layers: couches uniques après la boucle (décodage final)
+            num_recurrent_steps: nombre d'itérations T du bloc récurrent
+            use_lti_injection: active l'injection LTI h(t+1) = A·h(t) + B·e + block(h(t))
+            use_act_halting: active Adaptive Computation Time (arrêt dynamique par position)
+            act_halt_threshold: seuil de probabilité de halting pour ACT
+            depth_lora_rank: rang des adaptateurs LoRA par profondeur (0 = désactivé)
+    """
+    model_type = "minigpt"
+    
+    def __init__(
+        self,
+        vocab_size=32000,
+        block_size=256,
+        embed_dim=256,
+        depth=8,
+        heads=8,
+        num_kv_heads=None,
+        dropout=0.1,
+        hidden_dim=512,
+        weight_sharing="none",
+        use_rope=True,
+        use_gradient_checkpointing=False,
+        # Recurrent-Depth Transformer params
+        num_prelude_layers=2,
+        num_coda_layers=2,
+        num_recurrent_steps=8,
+        use_lti_injection=True,
+        use_act_halting=True,
+        act_halt_threshold=0.99,
+        depth_lora_rank=8,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.heads = heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else heads
+        self.dropout = dropout
+        self.hidden_dim = hidden_dim
+        self.weight_sharing = weight_sharing
+        self.use_rope = use_rope
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+
+        # RDT
+        self.num_prelude_layers = num_prelude_layers
+        self.num_coda_layers = num_coda_layers
+        self.num_recurrent_steps = num_recurrent_steps
+        self.use_lti_injection = use_lti_injection
+        self.use_act_halting = use_act_halting
+        self.act_halt_threshold = act_halt_threshold
+        self.depth_lora_rank = depth_lora_rank
+

model.safetensors

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

model_components.py

@@ -0,0 +1,458 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+from transformers import PreTrainedModel
+from .configuration_minigpt import MiniGPTConfig
+
+
+# ---------------------------------------------------------------------------
+# Recurrent-Depth Transformer (RDT) components
+# Ref: OpenMythos / Parcae (Prairie et al., 2026)
+# ---------------------------------------------------------------------------
+
+
+class LTIInjection(nn.Module):
+    """Linear Time-Invariant state injection for the recurrent loop.
+
+    At each iteration t:
+        h(t+1) = A · h(t) + B · e + transformer_output
+
+    Stability guarantee: ρ(A) < 1 by construction.
+    A is parameterised as  A = U · diag(σ(a)) · V^T  where  σ(a) ∈ (0, 1)
+    via a sigmoid on unconstrained parameters *a*.  Because the singular
+    values are all < 1, the spectral radius is guaranteed < 1.
+    We use a *diagonal* A for efficiency (embed_dim params, not embed_dim²).
+    """
+
+    def __init__(self, embed_dim: int):
+        super().__init__()
+        self.a_logit = nn.Parameter(torch.zeros(embed_dim))
+        self.b_scale = nn.Parameter(torch.ones(embed_dim) * 0.1)
+
+    def forward(self, h: torch.Tensor, encoded_input: torch.Tensor) -> torch.Tensor:
+        a = torch.sigmoid(self.a_logit).clamp(max=0.95)
+        return a * h + self.b_scale * encoded_input
+
+
+class DepthLoRA(nn.Module):
+    """Low-rank adaptation applied per recurrent depth step.
+
+    Adds a small  delta = (x @ A) @ B  to the hidden state, where A and B
+    have shape (embed_dim, rank) and (rank, embed_dim) respectively.
+    Each loop iteration t uses its own DepthLoRA so behavior differs per step.
+    """
+
+    def __init__(self, embed_dim: int, rank: int):
+        super().__init__()
+        self.down = nn.Linear(embed_dim, rank, bias=False)
+        self.up = nn.Linear(rank, embed_dim, bias=False)
+        nn.init.kaiming_uniform_(self.down.weight, a=math.sqrt(5))
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.up(self.down(x))
+
+
+class ACTHalting(nn.Module):
+    """Adaptive Computation Time halting mechanism.
+
+    Per-position scalar that learns when to stop the recurrent loop.
+    Produces a halting probability p(t) at each step t; once cumulative
+    probability exceeds *threshold*, the position stops receiving updates.
+
+    Returns:
+        weighted_state:  Σ_t  w(t) · h(t)  — the ACT-blended output
+        act_loss:        ponder cost  (Σ cumul_prob)  for regularisation
+    """
+
+    def __init__(self, embed_dim: int, threshold: float = 0.99):
+        super().__init__()
+        self.halt_proj = nn.Linear(embed_dim, 1, bias=True)
+        self.threshold = threshold
+
+    def compute_halt_prob(self, h: torch.Tensor) -> torch.Tensor:
+        """h: (B, T, D) -> (B, T, 1) halt probability."""
+        return torch.sigmoid(self.halt_proj(h))
+
+
+class RoPEEmbedding(nn.Module):
+    """Rotary Position Embedding (RoPE) comme utilisé dans LLaMA et autres LLMs modernes.
+    
+    RoPE encode les positions directement dans les queries et keys via des rotations,
+    sans nécessiter de paramètres apprenables.
+    
+    Args:
+        dim: Dimension de chaque tête d'attention (embed_dim // num_heads)
+        max_seq_len: Longueur de séquence maximale
+        base: Base pour le calcul des fréquences (10000 par défaut)
+    """
+    def __init__(self, dim, max_seq_len=2048, base=10000):
+        super().__init__()
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        self.base = base
+        
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        
+        t = torch.arange(max_seq_len).type_as(self.inv_freq)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
+        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
+    
+    def rotate_half(self, x):
+        """Rotation de moitié des dimensions."""
+        x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
+        return torch.cat((-x2, x1), dim=-1)
+    
+    def forward(self, q, k, offset=0):
+        """Applique RoPE aux queries et keys.
+        
+        Args:
+            q: queries [batch, heads, seq_len, head_dim]
+            k: keys [batch, heads, seq_len, head_dim]
+            offset: position offset for KV cache (number of previously cached tokens)
+        
+        Returns:
+            q_rot, k_rot: queries et keys avec positions encodées
+        """
+        seq_len = q.shape[2]
+        
+        cos = self.cos_cached[:, :, offset:offset + seq_len, :]
+        sin = self.sin_cached[:, :, offset:offset + seq_len, :]
+        
+        q_rot = (q * cos) + (self.rotate_half(q) * sin)
+        k_rot = (k * cos) + (self.rotate_half(k) * sin)
+        
+        return q_rot, k_rot
+
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization (Zhang & Sennrich, 2019).
+
+    Faster than LayerNorm: no mean-subtraction, no bias, single reduction.
+    Used by LLaMA, Mistral, Gemma, Qwen.
+    """
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU activation function as described in the Super Tiny LM paper.
+    SwiGLU(x) = (Swish(xW) ⊗ xV)W2
+    where Swish(x) = SiLU(x) in PyTorch
+    """
+    def __init__(self, embed_dim, hidden_dim):
+        super().__init__()
+        self.w = nn.Linear(embed_dim, hidden_dim, bias=False)
+        self.v = nn.Linear(embed_dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, embed_dim, bias=False)
+    
+    def forward(self, x):
+        return self.w2(F.silu(self.w(x)) * self.v(x))
+
+
+class SelfAttention(nn.Module):
+    """Multi-Head Attention with optional Grouped-Query Attention (GQA).
+
+    When num_kv_heads < num_heads, K and V projections are smaller and each
+    KV head is shared across num_heads // num_kv_heads query heads.
+    When num_kv_heads == num_heads this is standard MHA (fully backward-compatible).
+    """
+    def __init__(self, embed_dim, heads, dropout, max_seq_len=2048, use_rope=True, num_kv_heads=None):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else heads
+        self.head_dim = embed_dim // heads
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.use_rope = use_rope
+
+        assert heads % self.num_kv_heads == 0, (
+            f"num_heads ({heads}) must be divisible by num_kv_heads ({self.num_kv_heads})"
+        )
+
+        self.q_proj = nn.Linear(embed_dim, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(embed_dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.out = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.attn_dropout = dropout
+        self.resid_dropout = nn.Dropout(dropout)
+
+        if use_rope:
+            self.rope = RoPEEmbedding(self.head_dim, max_seq_len=max_seq_len)
+
+    def forward(self, x, mask=None, past_kv=None, use_cache=False):
+        B, T, C = x.size()
+
+        q = self.q_proj(x).reshape(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).reshape(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).reshape(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        if self.use_rope:
+            offset = past_kv[0].shape[2] if past_kv is not None else 0
+            q, k = self.rope(q, k, offset=offset)
+
+        if past_kv is not None:
+            k = torch.cat([past_kv[0], k], dim=2)
+            v = torch.cat([past_kv[1], v], dim=2)
+
+        new_kv = (k, v) if use_cache else None
+
+        # GQA: expand KV heads to match Q heads via repeat
+        k_expanded = k.repeat_interleave(self.num_kv_groups, dim=1) if self.num_kv_groups > 1 else k
+        v_expanded = v.repeat_interleave(self.num_kv_groups, dim=1) if self.num_kv_groups > 1 else v
+
+        is_causal = past_kv is None
+        attn = F.scaled_dot_product_attention(
+            q, k_expanded, v_expanded,
+            attn_mask=None,
+            is_causal=is_causal,
+            dropout_p=self.attn_dropout if self.training else 0.0,
+        )
+        attn = attn.transpose(1, 2).contiguous().view(B, T, C)
+        return self.resid_dropout(self.out(attn)), new_kv
+
+class TransformerBlock(nn.Module):
+    def __init__(self, embed_dim, heads, dropout=0.1, hidden_dim=512, layerdrop=0.1,
+                 shared_ff=None, max_seq_len=2048, use_rope=True, num_kv_heads=None):
+        super().__init__()
+        self.attn = SelfAttention(embed_dim, heads, dropout, max_seq_len=max_seq_len,
+                                  use_rope=use_rope, num_kv_heads=num_kv_heads)
+        self.ln1 = RMSNorm(embed_dim)
+        self.ff = shared_ff if shared_ff is not None else SwiGLU(embed_dim, hidden_dim)
+        self.ln2 = RMSNorm(embed_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.layerdrop = layerdrop
+
+    def forward(self, x, mask=None, past_kv=None, use_cache=False):
+        if self.training and torch.rand(1).item() < self.layerdrop:
+            return (x, None) if use_cache else x
+        attn_out, new_kv = self.attn(self.ln1(x), mask, past_kv=past_kv, use_cache=use_cache)
+        x = x + self.dropout(attn_out)
+        x = x + self.dropout(self.ff(self.ln2(x)))
+        return (x, new_kv) if use_cache else x
+    
+    def forward_checkpointed(self, x, mask=None):
+        """Version avec gradient checkpointing (training only, no cache)."""
+        attn_out, _ = self.attn(self.ln1(x), mask)
+        x = x + self.dropout(attn_out)
+        x = x + self.dropout(self.ff(self.ln2(x)))
+        return x
+
+class MiniGPT(PreTrainedModel):
+    config_class = MiniGPTConfig
+    
+    def __init__(self, config=None, **kwargs):
+        """
+        Initialise le modèle MiniGPT.
+        
+        Args:
+            config: Instance de MiniGPTConfig ou None. Si None, les paramètres
+                   doivent être fournis via kwargs.
+            **kwargs: Paramètres du modèle si config n'est pas fourni.
+        """
+        # Si config n'est pas fourni, créer une config à partir des kwargs
+        if config is None:
+            config = MiniGPTConfig(**kwargs)
+        
+        super().__init__(config)
+        
+        # Extraire les paramètres de la config
+        vocab_size = config.vocab_size
+        block_size = config.block_size
+        embed_dim = config.embed_dim
+        depth = config.depth
+        heads = config.heads
+        num_kv_heads = config.num_kv_heads
+        dropout = config.dropout
+        hidden_dim = config.hidden_dim
+        weight_sharing = config.weight_sharing
+        use_rope = config.use_rope
+        use_gradient_checkpointing = config.use_gradient_checkpointing
+        self.token_emb = nn.Embedding(vocab_size, embed_dim)
+        self.use_rope = use_rope
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        
+        # Positional embeddings uniquement si on n'utilise pas RoPE
+        if not use_rope:
+            self.pos_emb = nn.Embedding(block_size, embed_dim)
+        else:
+            self.pos_emb = None
+        
+        self.depth = depth
+        self.weight_sharing = weight_sharing
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.embed_dim = embed_dim
+        self.heads = heads
+        self.hidden_dim = hidden_dim
+        
+        # Créer les blocs selon le type de weight sharing
+        if weight_sharing == "none":
+            self.blocks = nn.ModuleList([
+                TransformerBlock(embed_dim, heads, dropout, hidden_dim, layerdrop=0.1,
+                                 max_seq_len=block_size, use_rope=use_rope,
+                                 num_kv_heads=num_kv_heads)
+                for _ in range(depth)
+            ])
+        elif weight_sharing == "ffn":
+            shared_ff = SwiGLU(embed_dim, hidden_dim)
+            self.blocks = nn.ModuleList([
+                TransformerBlock(embed_dim, heads, dropout, hidden_dim, layerdrop=0.1,
+                                 shared_ff=shared_ff, max_seq_len=block_size,
+                                 use_rope=use_rope, num_kv_heads=num_kv_heads)
+                for _ in range(depth)
+            ])
+        elif weight_sharing == "full":
+            self.shared_block = TransformerBlock(embed_dim, heads, dropout, hidden_dim, layerdrop=0.1,
+                                                 max_seq_len=block_size, use_rope=use_rope,
+                                                 num_kv_heads=num_kv_heads)
+            self.blocks = None
+        else:
+            raise ValueError(f"weight_sharing doit être 'none', 'ffn' ou 'full', pas '{weight_sharing}'")
+        
+        self.ln_f = RMSNorm(embed_dim)
+        self.head = nn.Linear(embed_dim, vocab_size, bias=False)
+        self.head.weight = self.token_emb.weight #On réutilise les poids de la matrice token_emb pour les tetes 
+        self.block_size = block_size
+        self.apply(self._init_weights)
+
+    def forward(self, idx):
+        B, T = idx.shape
+        
+        # Token embeddings
+        x = self.token_emb(idx)
+        
+        # Ajouter positional embeddings uniquement si on n'utilise pas RoPE
+        if not self.use_rope:
+            pos = torch.arange(0, T, device=idx.device).unsqueeze(0)
+            x = x + self.pos_emb(pos)
+
+        mask = torch.tril(torch.ones(T, T, device=idx.device)).unsqueeze(0).unsqueeze(0)
+        
+        # Gradient checkpointing : économise VRAM en recalculant les activations
+        if self.use_gradient_checkpointing and self.training:
+            if self.weight_sharing == "full":
+                for _ in range(self.depth):
+                    x = checkpoint(self.shared_block.forward_checkpointed, x, mask, use_reentrant=False)
+            else:
+                for block in self.blocks:
+                    x = checkpoint(block.forward_checkpointed, x, mask, use_reentrant=False)
+        else:
+            # Mode normal (pas de checkpointing)
+            if self.weight_sharing == "full":
+                for _ in range(self.depth):
+                    x = self.shared_block(x, mask)
+            else:
+                for block in self.blocks:
+                    x = block(x, mask)
+
+        x = self.ln_f(x)
+        return self.head(x)
+    
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, RMSNorm):
+            torch.nn.init.ones_(module.weight)
+            
+            
+    def count_parameters(self):
+        """Compte le nombre de paramètres selon le type de weight sharing et use_rope."""
+        total = sum(p.numel() for p in self.parameters())
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        
+        # Détails par composant
+        token_emb_params = self.token_emb.weight.numel()
+        pos_emb_params = self.pos_emb.weight.numel() if self.pos_emb is not None else 0
+        embedding_params = token_emb_params + pos_emb_params
+        
+        if self.weight_sharing == "full":
+            block_params = sum(p.numel() for p in self.shared_block.parameters())
+        else:
+            block_params = sum(p.numel() for p in self.blocks.parameters())
+        
+        return {
+            "total": total,
+            "trainable": trainable,
+            "embedding": embedding_params,
+            "token_emb": token_emb_params,
+            "pos_emb": pos_emb_params,
+            "blocks": block_params,
+            "head": 0,  # Head partage les poids avec embedding
+            "weight_sharing": self.weight_sharing,
+            "use_rope": self.use_rope
+        }
+    
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None, min_new_tokens=0, eos_token_id=None):
+        """
+        Génération de texte avec contrôle de la diversité.
+        
+        Args:
+            idx: Context initial [batch, seq_len]
+            max_new_tokens: Nombre de tokens à générer
+            temperature: Contrôle la diversité (0.1=conservateur, 1.0=normal, 2.0=créatif)
+            top_k: Garde seulement les k tokens les plus probables
+            top_p: Nucleus sampling, garde les tokens dont la somme des probas = p
+            min_new_tokens: Génère au moins ce nombre de tokens avant d'autoriser l'arrêt sur eos_token_id
+            eos_token_id: Id du token EOS pour stopper la génération (optionnel)
+        """
+        for step in range(max_new_tokens):
+            idx_cond = idx[:, -self.block_size:]
+            logits = self(idx_cond)
+            logits = logits[:, -1, :]
+            
+            # Appliquer la température
+            if temperature != 1.0:
+                logits = logits / temperature
+            
+            # Top-k filtering
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            
+            # Top-p (nucleus) filtering
+            if top_p is not None:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                
+                # Retirer les tokens au-delà du seuil top_p
+                sorted_indices_to_remove = cumulative_probs > top_p
+                # Garder au moins le premier token
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                sorted_indices_to_remove[..., 0] = 0
+                
+                # Scatter les valeurs -inf
+                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                logits[indices_to_remove] = -float('Inf')
+            
+            # Échantillonner
+            probs = F.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+
+            # Éviter un EOS trop tôt
+            if eos_token_id is not None and step < min_new_tokens:
+                while next_token.item() == eos_token_id:
+                    next_token = torch.multinomial(probs, num_samples=1)
+
+            idx = torch.cat((idx, next_token), dim=1)
+
+            # Arrêt précoce si EOS après le minimum requis
+            if eos_token_id is not None and step >= min_new_tokens and next_token.item() == eos_token_id:
+                break
+        return idx

modeling_minigpt.py

@@ -0,0 +1,209 @@
+"""
+Modèle MiniGPT pour Hugging Face Transformers.
+
+Ce fichier contient MiniGPTForCausalLM qui est la classe standard
+attendue par Hugging Face pour les modèles de génération de texte.
+
+MiniGPTForCausalLM hérite de MiniGPTModel et ajoute uniquement la tête de langage.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from .configuration_minigpt import MiniGPTConfig
+from .modeling_minigpt_core import MiniGPTModel
+
+
+
+class MiniGPTForCausalLM(PreTrainedModel, GenerationMixin):
+    """
+    MiniGPT model avec une tête de langage pour la génération de texte.
+    
+    Cette classe est compatible avec l'écosystème Hugging Face et peut être
+    utilisée avec AutoModelForCausalLM une fois enregistrée.
+    """
+    config_class = MiniGPTConfig
+    base_model_prefix = "model"
+    _tied_weights_keys = ["lm_head.weight"]
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MiniGPTModel(config)
+        self.lm_head = nn.Linear(config.embed_dim, config.vocab_size, bias=False)
+        self.post_init()
+
+    @property
+    def act_loss(self) -> torch.Tensor:
+        """ACT ponder cost from the recurrent loop (0 when not using RDT)."""
+        return self.model.act_loss
+
+    def set_inference_recurrent_steps(self, steps: int | None):
+        """Override the number of recurrent steps for inference.
+        Set to None to revert to training default."""
+        self.model.set_inference_recurrent_steps(steps)
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def tie_weights(self, **kwargs):
+        self.lm_head.weight = self.model.token_emb.weight
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is None:
+            raise ValueError("input_ids doit être fourni")
+
+        # Appel au modèle core
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        # Extraire les hidden states selon le format de retour
+        if return_dict:
+            hidden_states = outputs.last_hidden_state
+        else:
+            hidden_states = outputs[0]
+
+        # Appliquer la tête de langage
+        logits = self.lm_head(hidden_states)
+
+        # Calculer la loss si labels fournis
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1)
+            )
+
+        # Format de sortie selon return_dict
+        if not return_dict:
+            output = (logits,)
+            if loss is not None:
+                return (loss,) + output
+            return output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": True,
+        }
+    
+    @torch.no_grad()
+    def generate(self, input_ids=None, max_new_tokens=100, temperature=1.0, top_k=None, top_p=None,
+                 min_new_tokens=0, eos_token_id=None, repetition_penalty=1.0,
+                 no_repeat_ngram_size=0, **kwargs):
+        """
+        Génération de texte avec contrôle de la diversité et anti-répétition.
+        """
+        if input_ids is None:
+            raise ValueError("input_ids doit être fourni")
+
+        idx = input_ids
+        block_size = self.config.block_size
+        generated_ids: list[int] = []
+
+        for step in range(max_new_tokens):
+            idx_cond = idx[:, -block_size:]
+            logits = self.forward(idx_cond).logits
+            logits = logits[:, -1, :]
+
+            # Repetition penalty: reduce probability of tokens already seen
+            if repetition_penalty != 1.0:
+                seen = set(idx[0].tolist()) | set(generated_ids)
+                for tid in seen:
+                    if logits[0, tid] > 0:
+                        logits[0, tid] /= repetition_penalty
+                    else:
+                        logits[0, tid] *= repetition_penalty
+
+            # No-repeat n-gram: ban n-grams that already appeared
+            if no_repeat_ngram_size > 0 and len(generated_ids) >= no_repeat_ngram_size - 1:
+                all_ids = idx[0].tolist() + generated_ids
+                ngram_len = no_repeat_ngram_size
+                banned: set[int] = set()
+                suffix = tuple(all_ids[-(ngram_len - 1):])
+                for i in range(len(all_ids) - ngram_len + 1):
+                    if tuple(all_ids[i:i + ngram_len - 1]) == suffix:
+                        banned.add(all_ids[i + ngram_len - 1])
+                for tid in banned:
+                    logits[0, tid] = -float('Inf')
+
+            if temperature != 1.0:
+                logits = logits / temperature
+
+            if eos_token_id is not None and step < min_new_tokens:
+                logits[:, eos_token_id] = -float('Inf')
+
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+
+            if top_p is not None:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                sorted_indices_to_remove[..., 0] = 0
+                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                logits[indices_to_remove] = -float('Inf')
+
+            probs = F.softmax(logits, dim=-1)
+            if not torch.isfinite(probs).all() or (probs.sum(dim=-1) <= 0).any():
+                next_token = torch.argmax(logits, dim=-1, keepdim=True)
+            else:
+                next_token = torch.multinomial(probs, num_samples=1)
+
+            idx = torch.cat((idx, next_token), dim=1)
+            generated_ids.append(next_token.item())
+
+            if eos_token_id is not None and step >= min_new_tokens and next_token.item() == eos_token_id:
+                break
+
+        return idx
+    
+    def count_parameters(self):
+        """Délègue au modèle core MiniGPTModel."""
+        return self.model.count_parameters()

modeling_minigpt_core.py

@@ -0,0 +1,442 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+from transformers.modeling_outputs import BaseModelOutputWithPast
+
+from .model_components import (
+    RoPEEmbedding, RMSNorm, SwiGLU, SelfAttention, TransformerBlock,
+    LTIInjection, DepthLoRA, ACTHalting,
+)
+
+from .configuration_minigpt import MiniGPTConfig
+
+
+class MiniGPTModel(nn.Module):
+    """
+    Modèle core MiniGPT — sans tête LM, pure architecture Transformer.
+    NE DOIT PAS hériter de PreTrainedModel.
+
+    Supporte 4 modes via weight_sharing:
+        "none"              – couches indépendantes (classique)
+        "ffn"               – FFN partagé entre couches
+        "full"              – 1 bloc partagé appliqué depth fois
+        "recurrent_depth"   – Recurrent-Depth Transformer (Prelude → boucle → Coda)
+    """
+
+    def __init__(self, config: MiniGPTConfig):
+        super().__init__()
+
+        vocab_size = config.vocab_size
+        block_size = config.block_size
+        embed_dim = config.embed_dim
+        depth = config.depth
+        heads = config.heads
+        num_kv_heads = config.num_kv_heads
+        dropout = config.dropout
+        hidden_dim = config.hidden_dim
+
+        self.token_emb = nn.Embedding(vocab_size, embed_dim)
+        self.use_rope = config.use_rope
+        self.use_gradient_checkpointing = config.use_gradient_checkpointing
+        self.block_size = block_size
+        self.depth = depth
+        self.weight_sharing = config.weight_sharing
+
+        if not config.use_rope:
+            self.pos_emb = nn.Embedding(block_size, embed_dim)
+        else:
+            self.pos_emb = None
+
+        def _make_block(**extra):
+            return TransformerBlock(
+                embed_dim, heads, dropout, hidden_dim,
+                max_seq_len=block_size, use_rope=config.use_rope,
+                num_kv_heads=num_kv_heads, **extra,
+            )
+
+        # ----- Recurrent-Depth Transformer -----
+        if self.weight_sharing == "recurrent_depth":
+            self.num_prelude = config.num_prelude_layers
+            self.num_coda = config.num_coda_layers
+            self.num_recurrent_steps = config.num_recurrent_steps
+
+            self.prelude = nn.ModuleList([_make_block() for _ in range(self.num_prelude)])
+            self.recurrent_block = _make_block()
+            self.coda = nn.ModuleList([_make_block() for _ in range(self.num_coda)])
+
+            self.use_lti = config.use_lti_injection
+            if self.use_lti:
+                self.lti = LTIInjection(embed_dim)
+
+            self.depth_lora_rank = config.depth_lora_rank
+            if self.depth_lora_rank > 0:
+                self.depth_loras = nn.ModuleList([
+                    DepthLoRA(embed_dim, self.depth_lora_rank)
+                    for _ in range(self.num_recurrent_steps)
+                ])
+
+            self.use_act = config.use_act_halting
+            if self.use_act:
+                self.act = ACTHalting(embed_dim, threshold=config.act_halt_threshold)
+
+            self.blocks = None
+
+        # ----- Standard modes -----
+        elif self.weight_sharing == "none":
+            self.blocks = nn.ModuleList([_make_block() for _ in range(depth)])
+
+        elif self.weight_sharing == "ffn":
+            shared_ff = SwiGLU(embed_dim, hidden_dim)
+            self.blocks = nn.ModuleList([
+                _make_block(shared_ff=shared_ff) for _ in range(depth)
+            ])
+
+        elif self.weight_sharing == "full":
+            self.shared_block = _make_block()
+            self.blocks = None
+
+        else:
+            raise ValueError(
+                f"weight_sharing must be 'none', 'ffn', 'full' or 'recurrent_depth', "
+                f"got '{self.weight_sharing}'"
+            )
+
+        self.ln_f = RMSNorm(embed_dim)
+        self._act_loss = torch.tensor(0.0)
+        self._inference_recurrent_steps = None
+
+    # ------------------------------------------------------------------ #
+    #  Public API                                                         #
+    # ------------------------------------------------------------------ #
+
+    def get_input_embeddings(self):
+        return self.token_emb
+
+    def set_input_embeddings(self, value):
+        self.token_emb = value
+
+    def get_output_embeddings(self):
+        return None
+
+    @property
+    def act_loss(self) -> torch.Tensor:
+        """Ponder cost from ACT — add to training loss as regularisation."""
+        return self._act_loss
+
+    def set_inference_recurrent_steps(self, steps: int | None):
+        """Override the number of recurrent steps used at inference time.
+
+        Allows deeper reasoning without retraining. LoRA adapters are
+        reused cyclically for steps beyond the trained count.
+        Set to None to revert to the training default.
+        """
+        self._inference_recurrent_steps = steps
+
+    @property
+    def effective_recurrent_steps(self) -> int:
+        """Number of recurrent steps for the current forward pass."""
+        if self.weight_sharing != "recurrent_depth":
+            return 0
+        if not self.training and self._inference_recurrent_steps is not None:
+            return self._inference_recurrent_steps
+        return self.num_recurrent_steps
+
+    # ------------------------------------------------------------------ #
+    #  Forward                                                            #
+    # ------------------------------------------------------------------ #
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs
+    ):
+        return_dict = return_dict if return_dict is not None else True
+        use_cache = use_cache if use_cache is not None else False
+
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        if self.pos_emb is not None:
+            past_len = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+            pos = torch.arange(past_len, past_len + T, device=input_ids.device).unsqueeze(0)
+            x = x + self.pos_emb(pos)
+
+        mask = None
+        new_key_values = []
+
+        if self.weight_sharing == "recurrent_depth":
+            x = self._forward_recurrent_depth(
+                x, mask, past_key_values, use_cache, new_key_values,
+            )
+        elif self.use_gradient_checkpointing and self.training:
+            x = self._forward_checkpointed(x, mask)
+        elif use_cache:
+            x = self._forward_cached(x, mask, past_key_values, new_key_values)
+        else:
+            x = self._forward_standard(x, mask)
+
+        hidden_states = self.ln_f(x)
+        present_key_values = tuple(new_key_values) if use_cache else None
+
+        if not return_dict:
+            return (hidden_states,)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=present_key_values,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    # ------------------------------------------------------------------ #
+    #  RDT forward path                                                   #
+    # ------------------------------------------------------------------ #
+
+    def _forward_recurrent_depth(self, x, mask, past_key_values, use_cache, new_key_values):
+        """Prelude -> Recurrent loop (LTI + ACT + LoRA) -> Coda."""
+        kv_idx = 0
+
+        # --- Prelude ---
+        for block in self.prelude:
+            if self.use_gradient_checkpointing and self.training:
+                x = checkpoint(block.forward_checkpointed, x, mask, use_reentrant=False)
+            elif use_cache:
+                past_kv = past_key_values[kv_idx] if past_key_values is not None else None
+                x, kv = block(x, mask, past_kv=past_kv, use_cache=True)
+                new_key_values.append(kv)
+                kv_idx += 1
+            else:
+                x = block(x, mask)
+
+        encoded_input = x  # snapshot for LTI re-injection
+
+        # --- Recurrent loop ---
+        n_steps = self.effective_recurrent_steps
+        if self.use_act and not self.training:
+            x = self._recurrent_loop_act_inference(
+                x, encoded_input, mask, past_key_values, use_cache, new_key_values, kv_idx,
+                n_steps,
+            )
+            kv_idx += n_steps
+        else:
+            x, kv_idx = self._recurrent_loop(
+                x, encoded_input, mask, past_key_values, use_cache, new_key_values, kv_idx,
+                n_steps,
+            )
+
+        # --- Coda ---
+        for block in self.coda:
+            if self.use_gradient_checkpointing and self.training:
+                x = checkpoint(block.forward_checkpointed, x, mask, use_reentrant=False)
+            elif use_cache:
+                past_kv = past_key_values[kv_idx] if past_key_values is not None else None
+                x, kv = block(x, mask, past_kv=past_kv, use_cache=True)
+                new_key_values.append(kv)
+                kv_idx += 1
+            else:
+                x = block(x, mask)
+
+        return x
+
+    def _recurrent_loop(self, x, encoded_input, mask, past_key_values, use_cache, new_key_values, kv_idx,
+                         n_steps=None):
+        """Standard recurrent loop (training, or eval without ACT).
+
+        During training with ACT enabled, we compute halt probabilities and
+        accumulate the weighted output + ponder cost for the loss.
+        """
+        if n_steps is None:
+            n_steps = self.num_recurrent_steps
+        B, T, D = x.shape
+        device = x.device
+
+        use_act_train = self.use_act and self.training
+        if use_act_train:
+            cumul_prob = torch.zeros(B, T, 1, device=device)
+            remainder = torch.ones(B, T, 1, device=device)
+            weighted_state = torch.zeros_like(x)
+            still_running = torch.ones(B, T, 1, device=device, dtype=torch.bool)
+
+        for t in range(n_steps):
+            if self.use_gradient_checkpointing and self.training:
+                block_out = checkpoint(
+                    self.recurrent_block.forward_checkpointed, x, mask, use_reentrant=False,
+                )
+            elif use_cache:
+                past_kv = past_key_values[kv_idx] if past_key_values is not None else None
+                block_out, kv = self.recurrent_block(x, mask, past_kv=past_kv, use_cache=True)
+                new_key_values.append(kv)
+                kv_idx += 1
+            else:
+                block_out = self.recurrent_block(x, mask)
+
+            # LTI injection: h(t+1) = A·h(t) + B·e + block(h(t))
+            if self.use_lti:
+                x = self.lti(x, encoded_input) + block_out
+            else:
+                x = block_out
+
+            if self.depth_lora_rank > 0:
+                lora_idx = t % len(self.depth_loras)
+                x = x + self.depth_loras[lora_idx](x)
+
+            # ACT accumulation (training only)
+            if use_act_train:
+                p = self.act.compute_halt_prob(x)
+                still_running_f = still_running.float()
+                new_halted = (cumul_prob + p > self.act.threshold) & still_running
+                still_running = still_running & ~new_halted
+
+                w = torch.where(new_halted, remainder, p * still_running_f)
+                weighted_state = weighted_state + w * x
+                cumul_prob = cumul_prob + p * still_running_f
+                remainder = remainder - p * still_running_f
+
+        if use_act_train:
+            remainder_mask = still_running.float()
+            weighted_state = weighted_state + remainder * remainder_mask * x
+            self._act_loss = cumul_prob.sum() / (B * T)
+            x = weighted_state
+
+        return x, kv_idx
+
+    def _recurrent_loop_act_inference(
+        self, x, encoded_input, mask, past_key_values, use_cache, new_key_values, kv_idx,
+        n_steps=None,
+    ):
+        """Recurrent loop with early exit per position at inference time."""
+        if n_steps is None:
+            n_steps = self.num_recurrent_steps
+        B, T, D = x.shape
+        device = x.device
+
+        cumul_prob = torch.zeros(B, T, 1, device=device)
+        remainder = torch.ones(B, T, 1, device=device)
+        weighted_state = torch.zeros_like(x)
+        still_running = torch.ones(B, T, 1, device=device, dtype=torch.bool)
+
+        for t in range(n_steps):
+            if not still_running.any():
+                if use_cache:
+                    new_key_values.append(None)
+                continue
+
+            if use_cache:
+                past_kv = past_key_values[kv_idx + t] if past_key_values is not None else None
+                block_out, kv = self.recurrent_block(x, mask, past_kv=past_kv, use_cache=True)
+                new_key_values.append(kv)
+            else:
+                block_out = self.recurrent_block(x, mask)
+
+            if self.use_lti:
+                x = self.lti(x, encoded_input) + block_out
+            else:
+                x = block_out
+
+            if self.depth_lora_rank > 0:
+                lora_idx = t % len(self.depth_loras)
+                x = x + self.depth_loras[lora_idx](x)
+
+            p = self.act.compute_halt_prob(x)
+            still_running_f = still_running.float()
+            new_halted = (cumul_prob + p > self.act.threshold) & still_running
+            still_running = still_running & ~new_halted
+
+            w = torch.where(new_halted, remainder, p * still_running_f)
+            weighted_state = weighted_state + w * x
+            cumul_prob = cumul_prob + p * still_running_f
+            remainder = remainder - p * still_running_f
+
+        remainder_mask = still_running.float()
+        weighted_state = weighted_state + remainder * remainder_mask * x
+        self._act_loss = torch.tensor(0.0, device=device)
+
+        return weighted_state
+
+    # ------------------------------------------------------------------ #
+    #  Classic forward helpers (unchanged logic, just extracted)           #
+    # ------------------------------------------------------------------ #
+
+    def _forward_standard(self, x, mask):
+        if self.weight_sharing == "full":
+            for _ in range(self.depth):
+                x = self.shared_block(x, mask)
+        else:
+            for block in self.blocks:
+                x = block(x, mask)
+        return x
+
+    def _forward_cached(self, x, mask, past_key_values, new_key_values):
+        if self.weight_sharing == "full":
+            for layer_idx in range(self.depth):
+                past_kv = past_key_values[layer_idx] if past_key_values is not None else None
+                x, kv = self.shared_block(x, mask, past_kv=past_kv, use_cache=True)
+                new_key_values.append(kv)
+        else:
+            for layer_idx, block in enumerate(self.blocks):
+                past_kv = past_key_values[layer_idx] if past_key_values is not None else None
+                x, kv = block(x, mask, past_kv=past_kv, use_cache=True)
+                new_key_values.append(kv)
+        return x
+
+    def _forward_checkpointed(self, x, mask):
+        if self.weight_sharing == "full":
+            for _ in range(self.depth):
+                x = checkpoint(self.shared_block.forward_checkpointed, x, mask, use_reentrant=False)
+        else:
+            for block in self.blocks:
+                x = checkpoint(block.forward_checkpointed, x, mask, use_reentrant=False)
+        return x
+
+    # ------------------------------------------------------------------ #
+    #  Utilities                                                          #
+    # ------------------------------------------------------------------ #
+
+    def count_parameters(self):
+        """Compte le nombre de paramètres selon le mode."""
+        total = sum(p.numel() for p in self.parameters())
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+        token_emb_params = self.token_emb.weight.numel()
+        pos_emb_params = self.pos_emb.weight.numel() if self.pos_emb is not None else 0
+        embedding_params = token_emb_params + pos_emb_params
+
+        rdt_extra = 0
+        if self.weight_sharing == "recurrent_depth":
+            block_params = (
+                sum(p.numel() for p in self.prelude.parameters())
+                + sum(p.numel() for p in self.recurrent_block.parameters())
+                + sum(p.numel() for p in self.coda.parameters())
+            )
+            if self.use_lti:
+                rdt_extra += sum(p.numel() for p in self.lti.parameters())
+            if self.depth_lora_rank > 0:
+                rdt_extra += sum(p.numel() for p in self.depth_loras.parameters())
+            if self.use_act:
+                rdt_extra += sum(p.numel() for p in self.act.parameters())
+        elif self.weight_sharing == "full":
+            block_params = sum(p.numel() for p in self.shared_block.parameters())
+        else:
+            block_params = sum(p.numel() for p in self.blocks.parameters())
+
+        result = {
+            "total": total,
+            "trainable": trainable,
+            "embedding": embedding_params,
+            "token_emb": token_emb_params,
+            "pos_emb": pos_emb_params,
+            "blocks": block_params,
+            "head": 0,
+            "weight_sharing": self.weight_sharing,
+            "use_rope": self.use_rope,
+        }
+        if self.weight_sharing == "recurrent_depth":
+            result["rdt_extra"] = rdt_extra
+        return result

tokenizer_config.json

@@ -0,0 +1,21 @@
+{
+  "add_prefix_space": true,
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "extra_special_tokens": [
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<|end|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "CamembertTokenizer",
+  "unk_token": "<unk>"
+}