modeling_erk_linear.py

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

Yukleme (standart yol):
    from transformers import AutoModelForCausalLM, AutoTokenizer
    model = AutoModelForCausalLM.from_pretrained("ecloudtech/Erk-Linear", trust_remote_code=True)
    tok   = AutoTokenizer.from_pretrained("ecloudtech/Erk-Linear")

Gereksinimler: torch, transformers, flash-linear-attention, safetensors, huggingface_hub

Model Qwen3-14B mimarisine dayanir; 8 katmanin softmax dikkati subquadratic Gated DeltaNet
ile degistirilmis, kalan 32 katman softmax "cipa" olarak korunmustur. Govde agirliklari
`config.base_model` deposundan, GDN agirliklari bu depodan yuklenir; from_pretrained
ikisini birlestirip calisir bir nedensel dil modeli doner.
"""
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

try:                                                    # HF uzak kod (paket icinde)
    from .configuration_erk_linear import ErkLinearConfig
except ImportError:                                     # yerel kullanim (dosya yan yana)
    from configuration_erk_linear import ErkLinearConfig

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 _install_gdn(model, gdn_state, cfg, device, dtype):
    """Govde modelin secili self_attn katmanlarini GDN sarmalayicilariyla degistirir."""
    from fla.layers import GatedDeltaNet  # flash-linear-attention (triton -> GPU gerekir)

    layers = getattr(cfg, "gdn_layers", GDN_LAYERS)
    H = model.config.hidden_size
    for li in layers:
        gdn = GatedDeltaNet(
            hidden_size=H,
            head_dim=getattr(cfg, "gdn_head_dim", 128),
            num_heads=getattr(cfg, "gdn_num_heads", 40),
            use_gate=getattr(cfg, "gdn_use_gate", True),
            use_short_conv=getattr(cfg, "gdn_use_short_conv", True),
            mode=getattr(cfg, "gdn_mode", "chunk"),
        )
        prefix = f"L{li}."
        layer_sd = {k[len(prefix):]: v for k, v in gdn_state.items() if k.startswith(prefix)}
        if not layer_sd:
            raise ValueError(f"L{li} icin GDN agirligi bulunamadi ({cfg.gdn_weights_file})")
        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)
    return model


class ErkLinearForCausalLM(PreTrainedModel):
    """%20-lineer Erk hibridi.

    `from_pretrained`, govdeyi `config.base_model` deposundan yukler, bu depodaki GDN
    agirliklarini secili katmanlara takar ve elde edilen **calisir nedensel dil modelini**
    doner. Donen nesne standart bir transformers modelidir: `.generate()`, `.forward()`,
    `use_cache` ve chat sablonu oldugu gibi calisir.
    """

    config_class = ErkLinearConfig
    base_model_prefix = "erk_linear"

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path=None, *model_args, **kwargs):
        repo = pretrained_model_name_or_path or REPO_ID
        cfg = kwargs.pop("config", None)
        if not isinstance(cfg, ErkLinearConfig):
            cfg = ErkLinearConfig.from_pretrained(repo, **{
                k: kwargs[k] for k in ("revision", "token", "cache_dir") if k in kwargs
            })

        dtype = kwargs.pop("dtype", None) or kwargs.pop("torch_dtype", None) or torch.bfloat16
        device_map = kwargs.pop("device_map", None)
        kwargs.pop("trust_remote_code", None)

        base_kwargs = dict(kwargs)
        if device_map is not None:
            base_kwargs["device_map"] = device_map
        model = AutoModelForCausalLM.from_pretrained(
            cfg.base_model, dtype=dtype, trust_remote_code=True, **base_kwargs
        )
        if device_map is None:
            model = model.to("cuda" if torch.cuda.is_available() else "cpu")
        model.eval()

        gdn_path = hf_hub_download(
            repo_id=repo,
            filename=getattr(cfg, "gdn_weights_file", "gdn_weights.safetensors"),
            **{k: kwargs[k] for k in ("revision", "token", "cache_dir") if k in kwargs},
        )
        gdn_state = load_file(gdn_path)

        device = next(model.parameters()).device
        model = _install_gdn(model, gdn_state, cfg, device, dtype)
        model.config.erk_linear = {
            "gdn_layers": cfg.gdn_layers,
            "linear_ratio": cfg.linear_ratio,
            "base_model": cfg.base_model,
        }
        return model


def load_erk_linear(device="cuda", dtype=torch.bfloat16,
                    base_model=BASE_MODEL, repo_id=REPO_ID):
    """Geriye donuk uyumlu yardimci: (model, tokenizer) doner."""
    cfg = ErkLinearConfig(base_model=base_model)
    model = AutoModelForCausalLM.from_pretrained(base_model, dtype=dtype,
                                                 trust_remote_code=True).to(device).eval()
    gdn_path = hf_hub_download(repo_id=repo_id, filename=cfg.gdn_weights_file)
    model = _install_gdn(model, load_file(gdn_path), cfg, device, 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))