RubiRLM.py

"""Rubi-RLM: 1B-class Recursive Language Model (RLM) prototype.



Bu dosya, recursive düşünme + dual-loop öğrenme hedefiyle tasarlanmış bir

araştırma prototipi içerir.



Eklenen sohbet katmanı:

- İngilizce/Türkçe çift dilli chat şablonu

- HF tokenizer ile metin->id / id->metin köprüsü

- Tek mesaj veya interaktif chat CLI

"""

from __future__ import annotations

import argparse
import importlib
import importlib.util
from dataclasses import dataclass
from typing import List, Optional, Protocol, Sequence, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from rubi_train_stack import (
    TrainStackConfig,
    build_dataloader,
    build_dataset,
    build_optimizer,
    train_demo_steps,
)
from xqs_moe import build_deepspeed_moe
from xqs_stack import choose_moe_backend, detect_xqs_backends, format_backend_report
from x_quantum_sparse_ops import (
    build_linear,
    causal_scaled_dot_product_attention,
    fused_residual_add,
    maybe_compile_module,
    pack_rows,
    scatter_rows,
)


class TextTokenizer(Protocol):
    def encode(self, text: str, return_tensors: Optional[str] = None): ...

    def decode(self, token_ids: Sequence[int], skip_special_tokens: bool = True) -> str: ...


@dataclass
class ChatTurn:
    role: str
    content: str


@dataclass
class RLMConfig:
    vocab_size: int = 50_257
    max_seq_len: int = 2_048
    d_model: int = 2_048
    n_layers: int = 14
    n_heads: int = 16
    ff_mult: int = 4
    dropout: float = 0.1
    recurse_steps: int = 6
    critique_threshold: float = 0.20
    tie_embeddings: bool = True
    use_moe: bool = False
    moe_num_experts: int = 0
    moe_top_k: int = 2
    moe_expert_hidden: int = 0
    moe_router_jitter: float = 0.0
    moe_aux_loss_weight: float = 0.01
    use_layer_skip: bool = False
    layer_skip_threshold: float = 0.50
    layer_skip_target: float = 1.0
    layer_skip_aux_weight: float = 0.01
    use_ternary_weights: bool = False
    use_flash_attention: bool = False
    use_fused_ops: bool = False
    packed_execution: bool = False
    use_torch_compile: bool = False
    moe_backend: str = "auto"
    moe_ep_size: int = 1

    @classmethod
    def scale_1b(cls) -> "RLMConfig":
        return cls(
            vocab_size=50_257,
            max_seq_len=2_048,
            d_model=1_024,
            n_layers=10,
            n_heads=16,
            ff_mult=4,
            recurse_steps=6,
            critique_threshold=0.20,
            use_moe=True,
            moe_num_experts=32,
            moe_top_k=1,
            moe_expert_hidden=1_280,
            moe_router_jitter=0.01,
            moe_aux_loss_weight=0.01,
            use_layer_skip=True,
            layer_skip_threshold=0.80,
            layer_skip_target=0.03,
            layer_skip_aux_weight=0.01,
            use_ternary_weights=True,
            use_flash_attention=True,
            use_fused_ops=True,
            packed_execution=True,
            use_torch_compile=False,
            moe_backend="auto",
            moe_ep_size=1,
        )


class RMSNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-6):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
        return self.scale * (x / rms)


class DenseFeedForward(nn.Module):
    def __init__(self, cfg: RLMConfig):
        super().__init__()
        hidden = cfg.d_model * cfg.ff_mult
        self.up_proj = build_linear(cfg.d_model, hidden, ternary=cfg.use_ternary_weights)
        self.down_proj = build_linear(hidden, cfg.d_model, ternary=cfg.use_ternary_weights)
        self.dropout = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        return self.dropout(self.down_proj(F.gelu(self.up_proj(x)))), x.new_zeros(())


class FastSelfAttention(nn.Module):
    def __init__(self, cfg: RLMConfig):
        super().__init__()
        if cfg.d_model % cfg.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads.")
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.d_model // cfg.n_heads
        self.dropout = cfg.dropout
        self.use_flash_attention = cfg.use_flash_attention
        self.q_proj = build_linear(cfg.d_model, cfg.d_model, bias=False, ternary=cfg.use_ternary_weights)
        self.k_proj = build_linear(cfg.d_model, cfg.d_model, bias=False, ternary=cfg.use_ternary_weights)
        self.v_proj = build_linear(cfg.d_model, cfg.d_model, bias=False, ternary=cfg.use_ternary_weights)
        self.out_proj = build_linear(cfg.d_model, cfg.d_model, bias=False, ternary=cfg.use_ternary_weights)

    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        bsz, seq_len, _ = x.shape
        q = self.q_proj(x).view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        attn_out = causal_scaled_dot_product_attention(
            q,
            k,
            v,
            dropout_p=self.dropout,
            training=self.training,
        )
        attn_out = attn_out.transpose(1, 2).contiguous().view(bsz, seq_len, self.n_heads * self.head_dim)
        return self.out_proj(attn_out)


class MoEExpert(nn.Module):
    def __init__(self, d_model: int, hidden: int):
        super().__init__()
        self.up_proj = build_linear(d_model, hidden, ternary=True)
        self.down_proj = build_linear(hidden, d_model, ternary=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.gelu(self.up_proj(x)))


class MoEFeedForward(nn.Module):
    def __init__(self, cfg: RLMConfig):
        super().__init__()
        if cfg.moe_num_experts <= 0:
            raise ValueError("moe_num_experts must be positive when use_moe=True.")
        if cfg.moe_top_k <= 0 or cfg.moe_top_k > cfg.moe_num_experts:
            raise ValueError("moe_top_k must be in the range [1, moe_num_experts].")

        self.num_experts = cfg.moe_num_experts
        self.top_k = cfg.moe_top_k
        self.router_jitter = cfg.moe_router_jitter
        requested_backend = cfg.moe_backend.lower()
        self.backend = choose_moe_backend(prefer_deepspeed=requested_backend in {"auto", "deepspeed"}) if requested_backend != "native" else "native"
        self.router = build_linear(cfg.d_model, cfg.moe_num_experts, ternary=cfg.use_ternary_weights)
        self.experts = nn.ModuleList([MoEExpert(cfg.d_model, cfg.moe_expert_hidden) for _ in range(cfg.moe_num_experts)])
        self.deepspeed_moe = None
        if self.backend == "deepspeed":
            self.deepspeed_moe = build_deepspeed_moe(
                hidden_size=cfg.d_model,
                expert=MoEExpert(cfg.d_model, cfg.moe_expert_hidden),
                num_experts=cfg.moe_num_experts,
                top_k=cfg.moe_top_k,
                ep_size=cfg.moe_ep_size,
            )
            if self.deepspeed_moe is None:
                self.backend = "native"
        self.dropout = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        if self.deepspeed_moe is not None:
            out, aux_loss = self.deepspeed_moe(x)
            return self.dropout(out), aux_loss
        flat_x = x.reshape(-1, x.size(-1))
        router_logits = self.router(flat_x)
        if self.training and self.router_jitter > 0:
            router_logits = router_logits + torch.randn_like(router_logits) * self.router_jitter

        router_probs = F.softmax(router_logits, dim=-1)
        topk_weights, topk_indices = torch.topk(router_probs, self.top_k, dim=-1)
        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)

        mixed = flat_x.new_zeros(flat_x.shape)
        expert_load = router_probs.new_zeros(self.num_experts)

        for expert_id, expert in enumerate(self.experts):
            expert_mask = topk_indices == expert_id
            if not expert_mask.any():
                continue
            token_indices, slot_indices = expert_mask.nonzero(as_tuple=True)
            expert_inputs = flat_x.index_select(0, token_indices)
            expert_outputs = expert(expert_inputs)
            weights = topk_weights[token_indices, slot_indices].unsqueeze(-1)
            mixed.index_add_(0, token_indices, expert_outputs * weights)
            expert_load[expert_id] = float(token_indices.numel())

        mixed = self.dropout(mixed.view_as(x))
        importance = router_probs.mean(dim=0)
        load = expert_load / max(1, flat_x.size(0) * self.top_k)
        aux_loss = self.num_experts * torch.sum(importance * load)
        return mixed, aux_loss


class RecursiveBlock(nn.Module):
    def __init__(self, cfg: RLMConfig):
        super().__init__()

        self.use_layer_skip = cfg.use_layer_skip
        self.layer_skip_threshold = cfg.layer_skip_threshold
        self.layer_skip_target = cfg.layer_skip_target
        self.use_fused_ops = cfg.use_fused_ops
        self.packed_execution = cfg.packed_execution
        self.norm_attn = RMSNorm(cfg.d_model)
        self.norm_ff = RMSNorm(cfg.d_model)
        self.attn = FastSelfAttention(cfg)
        self.ffn = MoEFeedForward(cfg) if cfg.use_moe else DenseFeedForward(cfg)
        self.skip_router = build_linear(cfg.d_model, 1, ternary=cfg.use_ternary_weights) if cfg.use_layer_skip else None

        self.state_fuse = build_linear(cfg.d_model * 2, cfg.d_model, ternary=cfg.use_ternary_weights)
        self.state_update = build_linear(cfg.d_model, cfg.d_model, ternary=cfg.use_ternary_weights)
        self.state_gate = build_linear(cfg.d_model * 2, cfg.d_model, ternary=cfg.use_ternary_weights)

    def _run_core(

        self,

        x: torch.Tensor,

        state: torch.Tensor,

        attn_mask: Optional[torch.Tensor] = None,

    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        x_norm = self.norm_attn(x)
        attn_out = self.attn(x_norm, attn_mask=attn_mask)
        fuse_input = torch.cat([attn_out, state], dim=-1)
        gate = torch.sigmoid(self.state_gate(fuse_input))
        fused = self.state_fuse(fuse_input)
        fused = gate * fused + (1.0 - gate) * state
        if self.use_fused_ops:
            x = fused_residual_add(x, fused)
        else:
            x = x + fused
        ff_out, moe_aux_loss = self.ffn(self.norm_ff(x))
        if self.use_fused_ops:
            x = fused_residual_add(x, ff_out)
        else:
            x = x + ff_out
        new_state = torch.tanh(self.state_update(x))
        return x, new_state, moe_aux_loss

    def forward(

        self,

        x: torch.Tensor,

        state: torch.Tensor,

        attn_mask: Optional[torch.Tensor] = None,

    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        exec_prob = x.new_ones((x.size(0),))
        skip_aux_loss = x.new_zeros(())
        if self.skip_router is None:
            x, new_state, moe_aux_loss = self._run_core(x, state, attn_mask=attn_mask)
            return x, new_state, moe_aux_loss, skip_aux_loss, exec_prob.mean()

        router_input = x.mean(dim=1)
        exec_prob = torch.sigmoid(self.skip_router(router_input)).squeeze(-1)
        target = exec_prob.new_full(exec_prob.shape, self.layer_skip_target)
        skip_aux_loss = F.mse_loss(exec_prob, target)
        hard_gate = exec_prob >= self.layer_skip_threshold
        if not torch.any(hard_gate):
            return x, state, x.new_zeros(()), skip_aux_loss, exec_prob.mean()

        if torch.all(hard_gate):
            x_exec, state_exec, moe_aux_loss = self._run_core(x, state, attn_mask=attn_mask)
        elif self.packed_execution:
            active_indices = torch.nonzero(hard_gate, as_tuple=False).squeeze(-1)
            x_active, state_active = pack_rows(active_indices, x, state)
            x_active, state_active, moe_aux_loss = self._run_core(x_active, state_active, attn_mask=attn_mask)
            x_exec = scatter_rows(x, active_indices, x_active)
            state_exec = scatter_rows(state, active_indices, state_active)
        else:
            x_exec, state_exec, moe_aux_loss = self._run_core(x, state, attn_mask=attn_mask)

        if self.training:
            exec_gate = exec_prob + (hard_gate.to(exec_prob.dtype) - exec_prob).detach()
            exec_scale = exec_gate.view(-1, 1, 1)
            x_exec = x + exec_scale * (x_exec - x)
            state_exec = state + exec_scale * (state_exec - state)

        return x_exec, state_exec, moe_aux_loss, skip_aux_loss, exec_prob.mean()


class RubiRLM(nn.Module):
    def __init__(self, cfg: RLMConfig):
        super().__init__()
        self.cfg = cfg
        self._last_moe_aux_loss = torch.tensor(0.0)
        self._last_layer_skip_aux_loss = torch.tensor(0.0)

        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos_emb = nn.Embedding(cfg.max_seq_len, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)

        self.layers = nn.ModuleList([maybe_compile_module(RecursiveBlock(cfg), cfg.use_torch_compile) for _ in range(cfg.n_layers)])
        self.final_norm = RMSNorm(cfg.d_model)

        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.tok_emb.weight

        self.critique_head = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.d_model // 2),
            nn.GELU(),
            nn.Linear(cfg.d_model // 2, 1),
        )

    def _causal_mask(self, seq_len: int, device: torch.device) -> torch.Tensor:
        mask = torch.full((seq_len, seq_len), float("-inf"), device=device)
        return torch.triu(mask, diagonal=1)

    def _embed(self, input_ids: torch.Tensor) -> torch.Tensor:
        bsz, seq_len = input_ids.shape
        if seq_len > self.cfg.max_seq_len:
            raise ValueError(f"Girdi uzunluğu max_seq_len={self.cfg.max_seq_len} sınırını aşıyor.")
        pos = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(bsz, seq_len)
        return self.drop(self.tok_emb(input_ids) + self.pos_emb(pos))

    def forward_recursive(

        self,

        input_ids: torch.Tensor,

        steps: Optional[int] = None,

        stop_on_critique: bool = True,

        return_trace: bool = False,

    ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]:
        steps = steps or self.cfg.recurse_steps
        x = self._embed(input_ids)

        bsz, seq_len, d_model = x.shape
        states = [x.new_zeros((bsz, seq_len, d_model)) for _ in range(self.cfg.n_layers)]
        mask = self._causal_mask(seq_len, x.device)

        logits_trace: List[torch.Tensor] = []
        critique_trace: List[torch.Tensor] = []
        moe_aux_total = x.new_zeros(())
        layer_skip_aux_total = x.new_zeros(())

        for _ in range(steps):
            h = x
            new_states = []
            for layer, st in zip(self.layers, states):
                h, st_new, moe_aux, skip_aux, _ = layer(h, st, attn_mask=mask)
                new_states.append(st_new)
                moe_aux_total = moe_aux_total + moe_aux
                layer_skip_aux_total = layer_skip_aux_total + skip_aux
            states = new_states

            h_norm = self.final_norm(h)
            logits = self.lm_head(h_norm)
            pooled = h_norm[:, -1, :]
            critique = torch.sigmoid(self.critique_head(pooled)).squeeze(-1)

            logits_trace.append(logits)
            critique_trace.append(critique)
            x = h

            if stop_on_critique and torch.all(critique < self.cfg.critique_threshold):
                break

        denom = max(1, len(logits_trace) * len(self.layers))
        self._last_moe_aux_loss = moe_aux_total / denom
        self._last_layer_skip_aux_loss = layer_skip_aux_total / denom

        final_logits = logits_trace[-1]
        if return_trace:
            return final_logits, logits_trace, critique_trace
        return final_logits, [], critique_trace

    def training_loss(

        self,

        input_ids: torch.Tensor,

        target_ids: torch.Tensor,

        steps: Optional[int] = None,

        alpha_iterative: float = 0.30,

        beta_correction: float = 0.10,

    ) -> torch.Tensor:
        final_logits, trace, critique = self.forward_recursive(
            input_ids, steps=steps, stop_on_critique=False, return_trace=True
        )

        final_loss = F.cross_entropy(
            final_logits.view(-1, final_logits.size(-1)),
            target_ids.view(-1),
            ignore_index=-100,
        )

        if trace:
            iterative = 0.0
            for logits in trace[:-1]:
                iterative = iterative + F.cross_entropy(
                    logits.view(-1, logits.size(-1)),
                    target_ids.view(-1),
                    ignore_index=-100,
                )
            iterative = iterative / max(1, len(trace) - 1)
        else:
            iterative = final_loss.new_tensor(0.0)

        correction_bonus = 0.0
        if len(critique) > 1:
            start = critique[0].mean()
            end = critique[-1].mean()
            correction_bonus = torch.relu(end - start)

        total_loss = final_loss + alpha_iterative * iterative + beta_correction * correction_bonus
        if self.cfg.use_moe:
            total_loss = total_loss + self.cfg.moe_aux_loss_weight * self._last_moe_aux_loss
        if self.cfg.use_layer_skip:
            total_loss = total_loss + self.cfg.layer_skip_aux_weight * self._last_layer_skip_aux_loss
        return total_loss

    @torch.no_grad()
    def generate(

        self,

        input_ids: torch.Tensor,

        max_new_tokens: int = 64,

        temperature: float = 0.8,

        top_k: int = 50,

        steps: Optional[int] = None,

    ) -> torch.Tensor:
        self.eval()
        out = input_ids

        for _ in range(max_new_tokens):
            context = out[:, -self.cfg.max_seq_len :]
            logits, _, _ = self.forward_recursive(context, steps=steps, stop_on_critique=True, return_trace=False)
            next_logits = logits[:, -1, :] / max(temperature, 1e-5)

            if top_k > 0:
                values, _ = torch.topk(next_logits, min(top_k, next_logits.size(-1)))
                cutoff = values[:, [-1]]
                next_logits = torch.where(next_logits < cutoff, torch.full_like(next_logits, -1e9), next_logits)

            probs = F.softmax(next_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            out = torch.cat([out, next_token], dim=1)

        return out

    def generate_text(

        self,

        tokenizer: TextTokenizer,

        prompt: str,

        max_new_tokens: int = 128,

        temperature: float = 0.7,

        top_k: int = 50,

        steps: Optional[int] = None,

        device: Optional[torch.device] = None,

    ) -> str:
        device = device or next(self.parameters()).device
        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
        output_ids = self.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            steps=steps,
        )
        new_tokens = output_ids[0, input_ids.shape[1] :].tolist()
        return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

    def chat(

        self,

        tokenizer: TextTokenizer,

        history: List[ChatTurn],

        user_message: str,

        lang: str = "auto",

        max_new_tokens: int = 192,

        temperature: float = 0.7,

        top_k: int = 50,

        steps: Optional[int] = None,

        device: Optional[torch.device] = None,

    ) -> Tuple[str, List[ChatTurn]]:
        prompt = build_chat_prompt(history, user_message, lang=lang)
        assistant_reply = self.generate_text(
            tokenizer=tokenizer,
            prompt=prompt,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            steps=steps,
            device=device,
        )
        updated = history + [ChatTurn(role="user", content=user_message), ChatTurn(role="assistant", content=assistant_reply)]
        return assistant_reply, updated

    def outer_sleep_phase_step(

        self,

        optimizer: torch.optim.Optimizer,

        input_ids: torch.Tensor,

        target_ids: torch.Tensor,

        steps: Optional[int] = None,

    ) -> float:
        self.train()
        optimizer.zero_grad(set_to_none=True)
        loss = self.training_loss(input_ids, target_ids, steps=steps)
        loss.backward()
        nn.utils.clip_grad_norm_(self.parameters(), 1.0)
        optimizer.step()
        return float(loss.detach().item())


def estimate_parameters(cfg: RLMConfig) -> int:
    d = cfg.d_model
    total = cfg.vocab_size * d + cfg.max_seq_len * d
    attn_params = (4 * d * d) + (4 * d)
    state_params = (5 * d * d) + (3 * d)
    router_params = 0
    layer_skip_params = 0
    ff_params = (2 * d * d * cfg.ff_mult) + (d * cfg.ff_mult) + d
    if cfg.use_moe:
        router_params = (d * cfg.moe_num_experts) + cfg.moe_num_experts
        expert_params = (2 * d * cfg.moe_expert_hidden) + cfg.moe_expert_hidden + d
        ff_params = cfg.moe_num_experts * expert_params
    if cfg.use_layer_skip:
        layer_skip_params = d + 1
    per_layer = attn_params + state_params + router_params + layer_skip_params + ff_params + (2 * d)
    total += cfg.n_layers * per_layer
    total += d * (d // 2) + (d // 2) + (d // 2) + 1 + d
    if not cfg.tie_embeddings:
        total += d * cfg.vocab_size
    return total


def estimate_active_parameters(cfg: RLMConfig) -> int:
    d = cfg.d_model
    total = cfg.vocab_size * d + cfg.max_seq_len * d
    attn_params = (4 * d * d) + (4 * d)
    state_params = (5 * d * d) + (3 * d)
    router_params = 0
    layer_skip_params = 0
    ff_params = (2 * d * d * cfg.ff_mult) + (d * cfg.ff_mult) + d
    if cfg.use_moe:
        router_params = (d * cfg.moe_num_experts) + cfg.moe_num_experts
        expert_params = (2 * d * cfg.moe_expert_hidden) + cfg.moe_expert_hidden + d
        ff_params = cfg.moe_top_k * expert_params
    if cfg.use_layer_skip:
        layer_skip_params = d + 1
    routed_layer = attn_params + state_params + router_params + ff_params + (2 * d)
    routed_layer = cfg.layer_skip_target * routed_layer
    per_layer = layer_skip_params + routed_layer
    total += cfg.n_layers * per_layer
    total += d * (d // 2) + (d // 2) + (d // 2) + 1 + d
    if not cfg.tie_embeddings:
        total += d * cfg.vocab_size
    return int(total)


def language_system_prompt(lang: str) -> str:
    base = (
        "You are Rubi-RLM assistant. Reason step-by-step internally, be concise in final answer, "
        "self-correct if needed."
    )
    if lang == "tr":
        return base + " Yanıtlarını Türkçe ver."
    if lang == "en":
        return base + " Reply in English."
    return base + " Reply in the user's language (Turkish or English)."


def build_chat_prompt(history: List[ChatTurn], user_message: str, lang: str = "auto") -> str:
    lines = [f"<|system|>\n{language_system_prompt(lang)}"]
    for turn in history:
        role = "user" if turn.role.lower() == "user" else "assistant"
        lines.append(f"<|{role}|>\n{turn.content}")
    lines.append(f"\n{user_message}")
    lines.append("<|assistant|>\n")
    return "\n".join(lines)


def load_hf_tokenizer(tokenizer_name: str):
    if importlib.util.find_spec("transformers") is None:
        raise RuntimeError("transformers yüklü değil. `pip install transformers` ile kurun.")
    transformers = importlib.import_module("transformers")
    tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_name)
    if getattr(tokenizer, "pad_token_id", None) is None and getattr(tokenizer, "eos_token", None) is not None:
        tokenizer.pad_token = tokenizer.eos_token
    return tokenizer


def demo() -> None:
    cfg = RLMConfig(
        vocab_size=4096,
        max_seq_len=128,
        d_model=256,
        n_layers=4,
        n_heads=8,
        ff_mult=4,
        recurse_steps=4,
        use_moe=True,
        moe_num_experts=8,
        moe_top_k=2,
        moe_expert_hidden=384,
    )
    model = RubiRLM(cfg)
    x = torch.randint(0, cfg.vocab_size, (2, 32))
    y = torch.randint(0, cfg.vocab_size, (2, 32))

    loss = model.training_loss(x, y)
    print(f"demo_loss={loss.item():.4f}")

    out = model.generate(x[:, :8], max_new_tokens=8, steps=3)
    print("generated_shape=", tuple(out.shape))


def resolve_config(scale: str) -> RLMConfig:
    if scale == "1b":
        return RLMConfig.scale_1b()
    return RLMConfig(d_model=512, n_layers=8, n_heads=8, vocab_size=50_257, max_seq_len=512)


def runtime_torch_compile_available() -> bool:
    if not hasattr(torch, "compile"):
        return False
    if torch.cuda.is_available() and importlib.util.find_spec("triton") is None:
        return False
    return True


def apply_runtime_config_overrides(cfg: RLMConfig, args: argparse.Namespace) -> RLMConfig:
    cfg.moe_backend = getattr(args, "moe_backend", cfg.moe_backend)
    cfg.moe_ep_size = getattr(args, "moe_ep_size", cfg.moe_ep_size)
    requested_compile = bool(getattr(args, "use_torch_compile", cfg.use_torch_compile))
    cfg.use_torch_compile = requested_compile and runtime_torch_compile_available()
    return cfg


def maybe_load_checkpoint(model: RubiRLM, checkpoint: Optional[str], device: torch.device) -> None:
    if not checkpoint:
        return
    state = torch.load(checkpoint, map_location=device)
    if isinstance(state, dict) and "model_state_dict" in state:
        model.load_state_dict(state["model_state_dict"])
        return
    model.load_state_dict(state)


def run_single_chat(args: argparse.Namespace) -> None:
    cfg = apply_runtime_config_overrides(resolve_config(args.scale), args)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = RubiRLM(cfg).to(device)
    maybe_load_checkpoint(model, args.checkpoint, device)
    tokenizer = load_hf_tokenizer(args.tokenizer_name)

    history: List[ChatTurn] = []
    if args.interactive:
        print("Interactive chat başladı. Çıkmak için /exit yaz.")
        while True:
            user_msg = input("You> ").strip()
            if not user_msg:
                continue
            if user_msg.lower() in {"/exit", "exit", "quit"}:
                break
            reply, history = model.chat(
                tokenizer=tokenizer,
                history=history,
                user_message=user_msg,
                lang=args.lang,
                max_new_tokens=args.max_new_tokens,
                temperature=args.temperature,
                top_k=args.top_k,
                steps=args.steps,
                device=device,
            )
            print(f"Rubi> {reply}")
        return

    if not args.prompt:
        raise ValueError("--chat modunda --prompt veya --interactive gerekli.")

    reply, _ = model.chat(
        tokenizer=tokenizer,
        history=[],
        user_message=args.prompt,
        lang=args.lang,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_k=args.top_k,
        steps=args.steps,
        device=device,
    )
    print(reply)


def print_stack_report() -> None:
    report = detect_xqs_backends()
    print(format_backend_report(report))


def run_train_demo(args: argparse.Namespace) -> None:
    cfg = apply_runtime_config_overrides(resolve_config(args.scale), args)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = RubiRLM(cfg).to(device)
    maybe_load_checkpoint(model, args.checkpoint, device)

    train_cfg = TrainStackConfig(
        optimizer_name=args.optimizer_name,
        learning_rate=args.learning_rate,
        weight_decay=args.weight_decay,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=not args.disable_pin_memory,
        prefetch_factor=args.prefetch_factor,
        persistent_workers=not args.disable_persistent_workers,
        max_seq_len=cfg.max_seq_len,
        dataset_dir=args.dataset_dir,
        use_bf16=not args.disable_bf16,
    )
    dataset = build_dataset(
        dataset_dir=train_cfg.dataset_dir,
        vocab_size=cfg.vocab_size,
        max_seq_len=min(cfg.max_seq_len, args.train_seq_len),
        synthetic_samples=max(args.train_steps * args.batch_size * 2, 32),
    )
    dataloader = build_dataloader(dataset, train_cfg, shuffle=True)
    optimizer = build_optimizer(model, train_cfg)
    mean_loss, total_tokens = train_demo_steps(
        model=model,
        optimizer=optimizer,
        dataloader=dataloader,
        device=device,
        steps=args.train_steps,
        use_bf16=train_cfg.use_bf16,
    )
    print(
        f"train_demo optimizer={optimizer.__class__.__name__} steps={args.train_steps} "
        f"mean_loss={mean_loss:.4f} tokens={total_tokens:,} device={device}"
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Rubi-RLM recursive language model")
    parser.add_argument("--scale", choices=["1b", "tiny"], default="1b")
    parser.add_argument("--estimate-only", action="store_true")
    parser.add_argument("--demo", action="store_true")
    parser.add_argument("--train-demo", action="store_true")
    parser.add_argument("--stack-report", action="store_true")

    parser.add_argument("--chat", action="store_true", help="Türkçe/İngilizce sohbet modunu açar")
    parser.add_argument("--interactive", action="store_true", help="Interactive chat loop")
    parser.add_argument("--prompt", type=str, default="")
    parser.add_argument("--lang", choices=["auto", "tr", "en"], default="auto")
    parser.add_argument("--tokenizer-name", type=str, default="gpt2")
    parser.add_argument("--checkpoint", type=str, default=None)
    parser.add_argument("--steps", type=int, default=None)
    parser.add_argument("--max-new-tokens", type=int, default=192)
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--top-k", type=int, default=50)
    parser.add_argument("--optimizer-name", type=str, default="auto")
    parser.add_argument("--moe-backend", choices=["auto", "native", "deepspeed"], default="auto")
    parser.add_argument("--moe-ep-size", type=int, default=1)
    parser.add_argument("--use-torch-compile", action="store_true")
    parser.add_argument("--learning-rate", type=float, default=3e-4)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--prefetch-factor", type=int, default=4)
    parser.add_argument("--dataset-dir", type=str, default="")
    parser.add_argument("--train-steps", type=int, default=2)
    parser.add_argument("--train-seq-len", type=int, default=256)
    parser.add_argument("--disable-pin-memory", action="store_true")
    parser.add_argument("--disable-persistent-workers", action="store_true")
    parser.add_argument("--disable-bf16", action="store_true")
    args = parser.parse_args()

    if args.chat:
        run_single_chat(args)
        return

    if args.stack_report:
        print_stack_report()
        return

    if args.train_demo:
        run_train_demo(args)
        return

    if args.demo:
        demo()
        return

    cfg = apply_runtime_config_overrides(resolve_config(args.scale), args)
    n_params = estimate_parameters(cfg)
    active_params = estimate_active_parameters(cfg)
    print(f"Scale={args.scale}, estimated_params={n_params:,}, estimated_active_params={active_params:,}")
    if not args.estimate_only:
        model = RubiRLM(cfg)
        actual = sum(p.numel() for p in model.parameters())
        print(f"actual_params={actual:,}")


if __name__ == "__main__":
    main()