"""
Erk-Linear — Erk-14B'nin 8 dikkat katmanini Gated DeltaNet'e damitan %20-lineer hibrit.

Yukleme:
    # gerekli: pip install torch transformers flash-linear-attention safetensors huggingface_hub
    from modeling_erk_linear import load_erk_linear
    model, tokenizer = load_erk_linear()          # Erk-14B tabanini + GDN agirliklarini indirir
    out = model.generate(**tokenizer("Merhaba", return_tensors="pt").to(model.device))

Model, Qwen3-14B mimarisine dayanir; 8 katmanin softmax dikkati subquadratic Gated DeltaNet ile
degistirilmis, kalan 32 katman softmax "cipa" olarak korunmustur. Ayrinti: teknik rapor / GitHub.
"""
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

BASE_MODEL = "ecloudtech/Erk-14B"     # Qwen3-14B temelli Turkce model
REPO_ID = "ecloudtech/Erk-Linear"
GDN_LAYERS = [1, 3, 5, 7, 10, 36, 38, 39]   # %20 lineer, yayilmis yerlesim


class _GDNStateCache:
    """GatedDeltaNet'in get/update_layer_cache arayuzunun bekledigi minimal katman-durum tutucu.

    FLA'nin recurrent_state + conv_state'ini tek katman icin saklar; boylece cache'li uretim
    sirasinda GDN gecmis durumu adimlar arasi devreder.
    """
    def __init__(self):
        self._layers = []

    def __len__(self):
        return len(self._layers)

    def __getitem__(self, idx):
        return self._layers[idx]

    def update(self, layer_idx=0, recurrent_state=None, conv_state=None, **kwargs):
        while len(self._layers) <= layer_idx:
            self._layers.append({"recurrent_state": None, "conv_state": None})
        if recurrent_state is not None:
            self._layers[layer_idx]["recurrent_state"] = recurrent_state
        if conv_state is not None:
            self._layers[layer_idx]["conv_state"] = conv_state
        return self


class _GDNAttention(nn.Module):
    """Qwen3 self_attn cagri imzasiyla uyumlu Gated DeltaNet sarmalayici.

    Cache'li uretim (use_cache=True) sirasinda GDN'nin recurrent + convolution durumunu
    adimlar arasi devreder; boylece model.generate() ciktisi, tam-yeniden-hesaplama
    (use_cache=False) ile sayisal gurultuye kadar ayni olur. Referans amacli tek-dizi
    kullanim icindir (es zamanli/batch-paylasimli servis icin ayri durum yonetimi gerekir).
    """
    def __init__(self, gdn):
        super().__init__()
        gdn.layer_idx = 0
        self.gdn = gdn
        self._state = None

    def forward(self, hidden_states, *args, **kwargs):
        cache_position = kwargs.get("cache_position", None)
        seq_len = hidden_states.shape[1]
        new_sequence = (
            (cache_position is None and seq_len > 1)
            or (cache_position is not None and int(cache_position.reshape(-1)[0]) == 0)
        )
        if new_sequence or self._state is None:
            self._state = _GDNStateCache()
        out = self.gdn(hidden_states, use_cache=True, past_key_values=self._state)
        y = out[0] if isinstance(out, tuple) else out
        if isinstance(out, tuple) and len(out) >= 3 and out[2] is not None:
            self._state = out[2]
        return (y, None)


def load_erk_linear(device="cuda", dtype=torch.bfloat16,
                    base_model=BASE_MODEL, repo_id=REPO_ID):
    """Erk-Linear hibridini kurar ve (model, tokenizer) doner."""
    from fla.layers import GatedDeltaNet  # flash-linear-attention

    model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=dtype).to(device).eval()
    H = model.config.hidden_size

    gdn_path = hf_hub_download(repo_id=repo_id, filename="gdn_weights.safetensors")
    state = load_file(gdn_path)

    for li in GDN_LAYERS:
        gdn = GatedDeltaNet(hidden_size=H, head_dim=128, num_heads=40,
                            use_gate=True, use_short_conv=True, mode="chunk")
        prefix = f"L{li}."
        layer_sd = {k[len(prefix):]: v for k, v in state.items() if k.startswith(prefix)}
        gdn.load_state_dict(layer_sd)
        gdn = gdn.to(device).to(dtype).eval()
        model.model.layers[li].self_attn = _GDNAttention(gdn).to(device).to(dtype)

    tokenizer = AutoTokenizer.from_pretrained(base_model)
    return model, tokenizer


if __name__ == "__main__":
    m, t = load_erk_linear()
    ids = t("Türkiye'nin başkenti", return_tensors="pt").to(m.device)
    print(t.decode(m.generate(**ids, max_new_tokens=12)[0], skip_special_tokens=True))