Upload CrossDNA 28.6M pretrained files

28.6M/README.md

@@ -0,0 +1,56 @@
+---
+license: gpl-3.0
+---
+
+## Using CrossDNA 28.6M (140K sequence inputs)
+```python
+import os
+
+
+os.environ["DISABLE_TORCH_COMPILE"] = "1"
+os.environ["TORCHDYNAMO_DISABLE"] = "1"
+
+import torch
+
+
+if hasattr(torch, "compile"):
+    def _no_compile(fn=None, *args, **kwargs):
+        if fn is None:
+            def deco(f):
+                return f
+            return deco
+        return fn
+    torch.compile = _no_compile
+
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+MODEL_DIR = "/data/zhaol/projects/huggingface_crossdna_140K_len/crossdna"
+
+tok = AutoTokenizer.from_pretrained(
+    MODEL_DIR,
+    trust_remote_code=True,
+    local_files_only=True,
+)
+
+model = AutoModelForMaskedLM.from_pretrained(
+    MODEL_DIR,
+    trust_remote_code=True,
+    local_files_only=True,
+).eval()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+seq = "ACGT" * 128
+enc = tok(seq, return_tensors="pt", add_special_tokens=False)
+x = enc["input_ids"].to(device)
+
+with torch.inference_mode():
+    out = model(input_ids=x)
+    emb = model.extract_embeddings(x)
+
+print("input_ids.shape =", tuple(x.shape))
+print("logits.shape =", tuple(out.logits.shape))
+print("embeddings.shape =", tuple(emb.shape))
+
+```

28.6M/huggingface_crossdna_140K_len/crossdna/config.json

@@ -0,0 +1,106 @@
+{
+  "alphabet_size": 5,
+  "architectures": [
+    "CrossDNAForMaskedLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_crossdna.CrossDNAConfig",
+    "AutoModelForMaskedLM": "modeling_crossdna.CrossDNAForMaskedLM",
+    "AutoTokenizer": "tokenization_crossdna.CrossDNATokenizer"
+  },
+  "auto_update_ema_in_forward": true,
+  "aux_ce_weight": 0.0,
+  "block_size": 4096,
+  "bos_token_id": 2,
+  "bridge_dropout": 0.05,
+  "checkpoint_chunk_size": 1,
+  "checkpoint_core_layers": true,
+  "cls_token_id": 0,
+  "comba_cfg": {
+    "conv_size": 4,
+    "correction_factor": 0.02,
+    "expand_v": 1,
+    "head_dim": 64,
+    "hidden_size": 256,
+    "mode": "chunk",
+    "norm_eps": 1e-05,
+    "num_heads": 8,
+    "use_gate": true,
+    "use_short_conv": true
+  },
+  "compact_n_token_id": 4,
+  "core_checkpoint_chunk_size": 1,
+  "d_model": 256,
+  "depth": 8,
+  "detach_gate": false,
+  "disable_cross_view": false,
+  "dna_token_ids": {
+    "A": 7,
+    "C": 8,
+    "G": 9,
+    "N": 11,
+    "T": 10
+  },
+  "dna_token_start_id": 7,
+  "dna_tokens": [
+    "A",
+    "C",
+    "G",
+    "T",
+    "N"
+  ],
+  "drop_path_rates": [
+    0.0,
+    0.08
+  ],
+  "dropout": 0.1,
+  "dtype": "float32",
+  "ema_decay": 0.9995,
+  "eos_token_id": 1,
+  "for_representation": false,
+  "gate_freeze_steps": 1000,
+  "gate_sup_warmup_steps": 1000,
+  "gate_sup_weight": 0.003,
+  "gate_temp": 1.2,
+  "mask_token_id": 3,
+  "model_type": "crossdna",
+  "pad_token_id": 4,
+  "pretrain": true,
+  "rc_bidirectional_stopgrad": true,
+  "rc_max_weight": 0.2,
+  "rc_tau": 1.5,
+  "rc_warmup_steps": 2000,
+  "return_ab_logits": true,
+  "sem_max_weight": 0.1,
+  "sem_warmup_steps": 8000,
+  "sep_token_id": 1,
+  "streaming_loss": true,
+  "streaming_report_ab": true,
+  "transformer_cfg": {
+    "attn": {
+      "num_heads": 8,
+      "num_kv_heads": 8,
+      "qkv_bias": false,
+      "rope_theta": 10000,
+      "window_size": 512
+    },
+    "fuse_swiglu": true,
+    "hidden_act": "swish",
+    "hidden_ratio": 4.0,
+    "hidden_size": 256,
+    "max_position_embeddings": 4096,
+    "norm_eps": 1e-05
+  },
+  "transformers_version": "4.57.1",
+  "unk_token_id": 6,
+  "use_barlow": false,
+  "use_bridge": true,
+  "use_checkpointing": true,
+  "use_ema_teacher": true,
+  "use_final_conv": false,
+  "use_mem": false,
+  "use_rc_kl": false,
+  "use_s_scan": true,
+  "use_tv": false,
+  "vocab_size": 12
+}

28.6M/huggingface_crossdna_140K_len/crossdna/configuration_crossdna.py

@@ -0,0 +1,178 @@
+from transformers import PretrainedConfig
+
+
+class CrossDNAConfig(PretrainedConfig):
+    model_type = "crossdna"
+
+    def __init__(
+        self,
+        alphabet_size=5,
+        vocab_size=12,
+        dna_tokens=("A", "C", "G", "T", "N"),
+        dna_token_start_id=7,
+        compact_n_token_id=4,
+        dna_token_ids=None,
+
+        d_model=256,
+        block_size=4096,
+        depth=8,
+        drop_path_rates=(0.0, 0.08),
+        dropout=0.10,
+
+        pretrain=True,
+        use_s_scan=True,
+        for_representation=False,
+
+        use_bridge=True,
+        bridge_dropout=0.05,
+
+        use_ema_teacher=True,
+        ema_decay=0.9995,
+        auto_update_ema_in_forward=True,
+
+        use_mem=False,
+        use_rc_kl=False,
+        use_barlow=False,
+        use_tv=False,
+        use_final_conv=False,
+
+        sem_max_weight=0.10,
+        sem_warmup_steps=8000,
+        aux_ce_weight=0.0,
+
+        gate_freeze_steps=1000,
+        detach_gate=False,
+        gate_sup_weight=0.003,
+        gate_sup_warmup_steps=1000,
+        gate_temp=1.2,
+
+        use_checkpointing=True,
+        checkpoint_chunk_size=1,
+        checkpoint_core_layers=True,
+        core_checkpoint_chunk_size=1,
+
+        return_ab_logits=True,
+        streaming_loss=True,
+        streaming_report_ab=True,
+
+        disable_cross_view=False,
+
+        rc_bidirectional_stopgrad=True,
+        rc_max_weight=0.2,
+        rc_tau=1.5,
+        rc_warmup_steps=2000,
+
+        transformer_cfg=None,
+        comba_cfg=None,
+
+        pad_token_id=4,
+        bos_token_id=2,
+        eos_token_id=1,
+        sep_token_id=1,
+        cls_token_id=0,
+        mask_token_id=3,
+        unk_token_id=6,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            sep_token_id=sep_token_id,
+            cls_token_id=cls_token_id,
+            mask_token_id=mask_token_id,
+            unk_token_id=unk_token_id,
+            **kwargs,
+        )
+
+        self.alphabet_size = int(alphabet_size)
+        self.vocab_size = int(vocab_size)
+
+        self.dna_tokens = list(dna_tokens)
+        self.dna_token_start_id = int(dna_token_start_id)
+        self.compact_n_token_id = int(compact_n_token_id)
+
+        if dna_token_ids is None:
+            dna_token_ids = {
+                ch: self.dna_token_start_id + i
+                for i, ch in enumerate(self.dna_tokens)
+            }
+        self.dna_token_ids = dict(dna_token_ids)
+
+        self.d_model = int(d_model)
+        self.block_size = int(block_size)
+        self.depth = int(depth)
+        self.drop_path_rates = list(drop_path_rates) if drop_path_rates is not None else None
+        self.dropout = float(dropout)
+
+        self.pretrain = bool(pretrain)
+        self.use_s_scan = bool(use_s_scan)
+        self.for_representation = bool(for_representation)
+
+        self.use_bridge = bool(use_bridge)
+        self.bridge_dropout = float(bridge_dropout)
+
+        self.use_ema_teacher = bool(use_ema_teacher)
+        self.ema_decay = float(ema_decay)
+        self.auto_update_ema_in_forward = bool(auto_update_ema_in_forward)
+
+        self.use_mem = bool(use_mem)
+        self.use_rc_kl = bool(use_rc_kl)
+        self.use_barlow = bool(use_barlow)
+        self.use_tv = bool(use_tv)
+        self.use_final_conv = bool(use_final_conv)
+
+        self.sem_max_weight = float(sem_max_weight)
+        self.sem_warmup_steps = int(sem_warmup_steps)
+        self.aux_ce_weight = float(aux_ce_weight)
+
+        self.gate_freeze_steps = int(gate_freeze_steps)
+        self.detach_gate = bool(detach_gate)
+        self.gate_sup_weight = float(gate_sup_weight)
+        self.gate_sup_warmup_steps = int(gate_sup_warmup_steps)
+        self.gate_temp = float(gate_temp)
+
+        self.use_checkpointing = bool(use_checkpointing)
+        self.checkpoint_chunk_size = int(checkpoint_chunk_size)
+        self.checkpoint_core_layers = bool(checkpoint_core_layers)
+        self.core_checkpoint_chunk_size = int(core_checkpoint_chunk_size)
+
+        self.return_ab_logits = bool(return_ab_logits)
+        self.streaming_loss = bool(streaming_loss)
+        self.streaming_report_ab = bool(streaming_report_ab)
+
+        self.disable_cross_view = bool(disable_cross_view)
+
+        self.rc_bidirectional_stopgrad = bool(rc_bidirectional_stopgrad)
+        self.rc_max_weight = float(rc_max_weight)
+        self.rc_tau = float(rc_tau)
+        self.rc_warmup_steps = int(rc_warmup_steps)
+
+        self.transformer_cfg = transformer_cfg or {
+            "hidden_size": self.d_model,
+            "norm_eps": 1e-5,
+            "max_position_embeddings": self.block_size,
+            "hidden_ratio": 4.0,
+            "hidden_act": "swish",
+            "fuse_swiglu": True,
+            "attn": {
+                "num_heads": 8,
+                "num_kv_heads": 8,
+                "qkv_bias": False,
+                "window_size": 512,
+                "rope_theta": 10000,
+            },
+        }
+
+        self.comba_cfg = comba_cfg or {
+            "hidden_size": self.d_model,
+            "expand_v": 1,
+            "head_dim": 64,
+            "num_heads": 8,
+            "use_gate": True,
+            "mode": "chunk",
+            "use_short_conv": True,
+            "correction_factor": 0.02,
+            "conv_size": 4,
+            "norm_eps": 1e-5,
+        }

28.6M/huggingface_crossdna_140K_len/crossdna/last.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6788405dcb78275faeaf1b94ae14a27465a37a3ff1872b34bffdb16b07fdf333
+size 454953820

28.6M/huggingface_crossdna_140K_len/crossdna/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2bfbc40499d6b55d3ef4a3ba1e2f598dce843f083ad8e413621a7378b564a25a
+size 225803888

28.6M/huggingface_crossdna_140K_len/crossdna/modeling_crossdna.py

@@ -0,0 +1,1702 @@
+
+# -*- coding: utf-8 -*-
+import math
+import copy
+from functools import partial
+from contextlib import contextmanager
+from collections import namedtuple
+from typing import Dict, Optional, Tuple, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import MaskedLMOutput
+
+# speed
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+from fla.layers import comba
+from fla.layers.attn import Attention
+from fla.modules import GatedMLP as SambaMLP
+from fla.modules import RMSNorm
+
+try:
+    from omegaconf import OmegaConf
+except Exception:
+    OmegaConf = None
+
+try:
+    from .configuration_crossdna import CrossDNAConfig
+except ImportError:
+    from configuration_crossdna import CrossDNAConfig
+
+
+# ========================
+# OmegaConf helpers
+# ========================
+def _to_plain_container(x: Any) -> Any:
+    if OmegaConf is not None:
+        try:
+            if OmegaConf.is_config(x):
+                return OmegaConf.to_container(x, resolve=True)
+        except Exception:
+            pass
+    return x
+
+
+def _cfg_get(cfg: Any, key: str, default: Any = None) -> Any:
+    if cfg is None:
+        return default
+
+    try:
+        if isinstance(cfg, dict):
+            return cfg.get(key, default)
+    except Exception:
+        pass
+
+    try:
+        if hasattr(cfg, key):
+            return getattr(cfg, key)
+    except Exception:
+        pass
+
+    try:
+        return cfg[key]
+    except Exception:
+        return default
+
+
+# ========================
+# Utils
+# ========================
+def complement(seq: torch.Tensor) -> torch.Tensor:
+    """
+    compact DNA ids only:
+        A=0, C=1, G=2, T=3, N=4
+    """
+    perm = torch.tensor([3, 2, 1, 0, 4], device=seq.device, dtype=torch.long)
+    return perm[seq.long()].to(seq.dtype)
+
+
+def reverse_complement(seq: torch.Tensor) -> torch.Tensor:
+    comp = complement(seq)
+    return torch.flip(comp, dims=[1])
+
+
+def make_complement_perm(C=5, device=None, dtype=torch.float32):
+    perm = torch.arange(C, device=device)
+    if C >= 4:
+        perm[0] = 3
+        perm[1] = 2
+        perm[2] = 1
+        perm[3] = 0
+    if C >= 5:
+        perm[4] = 4
+
+    P = torch.zeros(C, C, device=device, dtype=dtype)
+    P[torch.arange(C, device=device), perm] = 1.0
+    return P, perm
+
+
+def ensure_finite(x: torch.Tensor, name: str):
+    if not torch.isfinite(x).all():
+        raise FloatingPointError(f"Non-finite values detected in {name}")
+    return x
+
+
+def linear_warmup_weight(step: int, warmup_steps: int, max_w: float):
+    if warmup_steps <= 0:
+        return max_w
+    if step <= 0:
+        return 0.0
+    if step >= warmup_steps:
+        return max_w
+    return max_w * (step / warmup_steps)
+
+
+def preferred_amp_dtype():
+    try:
+        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
+            return torch.bfloat16
+    except Exception:
+        pass
+    return torch.float16
+
+
+def one_hot_float(x: torch.Tensor, num_classes: int, *, dtype: torch.dtype) -> torch.Tensor:
+    B, L = x.shape
+    out = torch.zeros((B, L, num_classes), device=x.device, dtype=dtype)
+    out.scatter_(2, x.unsqueeze(-1), 1.0)
+    return out
+
+
+# ========================
+# RC / Barlow / TV
+# ========================
+def rc_consistency_kl(logits_A, logits_B_fwd, P, tau: float = 1.0, eps: float = 1e-6):
+    zA = logits_A.float() / tau
+    zB = logits_B_fwd.float() / tau
+    pA = F.softmax(zA, dim=-1)
+    logpA = F.log_softmax(zA, dim=-1)
+    pB = F.softmax(zB, dim=-1)
+    pB_comp = torch.matmul(pB, P.t()).clamp_min(eps)
+    logpB_comp = pB_comp.log()
+    kl = (pA * (logpA - logpB_comp)).sum(dim=-1).mean()
+    return kl * (tau * tau)
+
+
+def rc_consistency_bidirectional_stopgrad(logits_A, logits_B_fwd, P, tau: float = 1.5, eps: float = 1e-6):
+    zA = logits_A.float() / tau
+    zB = logits_B_fwd.float() / tau
+    with torch.no_grad():
+        pB_t = torch.matmul(F.softmax(zB, dim=-1), P.t()).clamp_min(eps)
+        logpB_t = pB_t.log()
+    loss_A = F.kl_div(F.log_softmax(zA, dim=-1), logpB_t, reduction="batchmean", log_target=True)
+
+    with torch.no_grad():
+        pA_t = torch.matmul(F.softmax(zA, dim=-1), P.t()).clamp_min(eps)
+        logpA_t = pA_t.log()
+    loss_B = F.kl_div(F.log_softmax(zB, dim=-1), logpA_t, reduction="batchmean", log_target=True)
+    return 0.5 * (tau * tau) * (loss_A + loss_B)
+
+
+def barlow_strand_loss_v2(z1, z2, λ_off=0.04, λ_diag=0.04, eps=1e-3):
+    B, L, H = z1.shape
+    n = B * L
+    z1 = z1.reshape(n, H)
+    z2 = z2.reshape(n, H)
+
+    def _std(z):
+        var = z.var(dim=0, unbiased=False)
+        return torch.sqrt(var + eps)
+
+    std1, std2 = _std(z1), _std(z2)
+    var_term = (F.relu(1 - std1).pow(2).mean() + F.relu(1 - std2).pow(2).mean())
+
+    z1 = (z1 - z1.mean(0)) / (std1 + eps)
+    z2 = (z2 - z2.mean(0)) / (std2 + eps)
+    c = (z1.t() @ z2) / max(1, n)
+    diag = torch.diagonal(c)
+    off = c - torch.diag_embed(diag)
+    cov = λ_diag * (1 - diag).pow(2).mean() + λ_off * off.pow(2).mean()
+    return var_term + cov
+
+
+def tv_mixed(h: torch.Tensor):
+    d1 = h[:, 1:, :] - h[:, :-1, :]
+    d2 = d1[:, 1:, :] - d1[:, :-1, :]
+    return d1.abs().mean() + d2.pow(2).mean()
+
+
+class Mlp(nn.Module):
+    def __init__(self, input_dimension, hidden_dimension=None, output_dimension=None,
+                 activation=F.gelu, return_residual=False):
+        super().__init__()
+        self.return_residual = return_residual
+        hd = hidden_dimension or input_dimension
+        od = output_dimension or input_dimension
+        self.linear1 = nn.Linear(input_dimension, hd)
+        self.activation = activation
+        self.linear2 = nn.Linear(hd, od)
+
+    def forward(self, x: torch.Tensor):
+        h = self.activation(self.linear1(x))
+        y = self.linear2(h)
+        return (y, x) if self.return_residual else y
+
+
+def create_comba_cls(comba_kwargs=None, device=None, dtype=None):
+    factory_kwargs = {}
+    if device is not None:
+        factory_kwargs["device"] = device
+    if dtype is not None:
+        factory_kwargs["dtype"] = dtype
+    try:
+        base_kwargs = dict(comba_kwargs or {})
+        mixer_cls = partial(comba.Comba, **base_kwargs, **factory_kwargs)
+    except ImportError:
+        class FallbackComba(nn.Module):
+            def forward(self, x, *args, **kwargs):
+                return x
+        mixer_cls = lambda *args, **kwargs: FallbackComba()
+    return mixer_cls
+
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, config: Any):
+        super().__init__()
+        config = _to_plain_container(config)
+
+        hidden_size = _cfg_get(config, "hidden_size")
+        norm_eps = _cfg_get(config, "norm_eps", 1e-5)
+        attn_cfg = _cfg_get(config, "attn", {}) or {}
+        attn_cfg = _to_plain_container(attn_cfg)
+
+        self.mixer_norm = RMSNorm(hidden_size=hidden_size, eps=norm_eps)
+        self.mixer = Attention(
+            hidden_size=hidden_size,
+            num_heads=_cfg_get(attn_cfg, "num_heads"),
+            num_kv_heads=_cfg_get(attn_cfg, "num_kv_heads"),
+            qkv_bias=_cfg_get(attn_cfg, "qkv_bias"),
+            window_size=_cfg_get(attn_cfg, "window_size"),
+            rope_theta=_cfg_get(attn_cfg, "rope_theta"),
+            max_position_embeddings=_cfg_get(config, "max_position_embeddings"),
+        )
+
+        self.mlp_norm = RMSNorm(hidden_size, eps=norm_eps)
+        self.mlp = SambaMLP(
+            hidden_size=hidden_size,
+            hidden_ratio=_cfg_get(config, "hidden_ratio", 4.0),
+            hidden_act=_cfg_get(config, "hidden_act", "swish"),
+            fuse_swiglu=_cfg_get(config, "fuse_swiglu", True),
+        )
+        self.pre_scale = 1.0 / math.sqrt(2.0)
+
+    def forward(self, hidden_states: torch.Tensor, cache_params: Optional[Any] = None, **kwargs) -> Tuple[torch.Tensor, Any]:
+        residual = hidden_states
+        x = self.mixer_norm(hidden_states)
+
+        amp_dtype = preferred_amp_dtype()
+        device_type = x.device.type if x.device.type in ["cuda", "cpu", "xpu"] else "cuda"
+
+        with torch.autocast(device_type=device_type, enabled=(device_type == "cuda"), dtype=amp_dtype):
+            x_scaled = x * self.pre_scale
+            attn_out, _, cache_params = self.mixer(hidden_states=x_scaled, past_key_values=cache_params, **kwargs)
+            attn_out = attn_out / self.pre_scale
+
+        ensure_finite(attn_out, "attention_out")
+        h = residual + attn_out.to(x.dtype)
+
+        residual = h
+        x = self.mlp_norm(h)
+        with torch.autocast(device_type=device_type, enabled=(device_type == "cuda"), dtype=amp_dtype):
+            x = self.mlp(x, **kwargs)
+
+        h = residual + x
+        ensure_finite(h, "block_output")
+        return h, cache_params
+
+
+class EnhancedHybridCore(nn.Module):
+    def __init__(self, hidden_dim, comba_cfg, transformer_cfg, layer_idx=0, device=None, dtype=None):
+        super().__init__()
+        comba_cfg = _to_plain_container(comba_cfg)
+        transformer_cfg = _to_plain_container(transformer_cfg)
+
+        self.comba_cls = create_comba_cls(comba_kwargs=comba_cfg, device=device, dtype=dtype)
+        try:
+            self.comba = self.comba_cls(layer_idx=layer_idx)
+        except TypeError:
+            self.comba = self.comba_cls()
+
+        self.transformer = SlidingWindowAttention(config=transformer_cfg)
+        self.gate = nn.Linear(hidden_dim * 2, hidden_dim)
+        self.out_norm = nn.LayerNorm(hidden_dim)
+
+    @staticmethod
+    def _first(x):
+        return x[0] if isinstance(x, tuple) else x
+
+    def forward(self, x):
+        orig_dtype = x.dtype
+        x_fp32 = x.float()
+        device_type = x.device.type if x.device.type in ["cuda", "cpu", "xpu"] else "cuda"
+
+        with torch.autocast(device_type=device_type, enabled=False):
+            m_out = self._first(self.comba(x_fp32))
+
+        m_out = m_out.to(orig_dtype)
+        del x_fp32
+
+        t_out, _ = self.transformer(m_out)
+
+        concat = torch.cat([m_out, t_out], dim=-1)
+        g = torch.sigmoid(self.gate(concat))
+        fused = g * t_out + (1 - g) * m_out
+        y = self.out_norm(fused)
+        ensure_finite(y, "EnhancedHybridCore.out")
+        return y
+
+
+class DeepEnhancedBranch(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int,
+        comba_cfg: Dict | None,
+        transformer_cfg: Any,
+        depth: int = 4,
+        drop_path_rates=None,
+        *,
+        device=None,
+        dtype=None,
+        checkpoint_core_layers: bool = False,
+        core_checkpoint_chunk_size: int = 1,
+    ):
+        super().__init__()
+        self.layers = nn.ModuleList()
+
+        transformer_cfg = _to_plain_container(transformer_cfg)
+        comba_cfg = _to_plain_container(comba_cfg)
+
+        self.checkpoint_core_layers = bool(checkpoint_core_layers)
+        self.core_checkpoint_chunk_size = int(core_checkpoint_chunk_size)
+
+        if drop_path_rates is None:
+            rates = [0.05 * (i / max(1, depth - 1)) for i in range(depth)]
+        elif isinstance(drop_path_rates, (float, int)):
+            rates = [float(drop_path_rates)] * depth
+        else:
+            dpr = list(_to_plain_container(drop_path_rates))
+            rates = dpr + [dpr[-1]] * (depth - len(dpr))
+
+        for i in range(depth):
+            layer_cfg = dict(transformer_cfg) if isinstance(transformer_cfg, dict) else transformer_cfg.copy()
+            layer_cfg["drop_path_prob"] = rates[i]
+            self.layers.append(EnhancedHybridCore(hidden_dim, comba_cfg, layer_cfg, i, device, dtype))
+
+        self.output_norm = nn.LayerNorm(hidden_dim)
+
+    def _run_layers(self, x: torch.Tensor, start: int, end: int):
+        out = x
+        for i in range(start, end):
+            out = self.layers[i](out)
+        return out
+
+    def forward(self, x: torch.Tensor):
+        if self.training and self.checkpoint_core_layers:
+            chunk = max(1, self.core_checkpoint_chunk_size)
+            for s in range(0, len(self.layers), chunk):
+                e = min(s + chunk, len(self.layers))
+
+                def _seg(inp, s=s, e=e):
+                    return self._run_layers(inp, s, e)
+
+                x = cp.checkpoint(_seg, x, use_reentrant=False)
+        else:
+            for layer in self.layers:
+                x = layer(x)
+
+        y = self.output_norm(x)
+        ensure_finite(y, "DeepEnhancedBranch.out")
+        return y
+
+
+class TokenBridge(nn.Module):
+    def __init__(self, hidden_dim: int, dropout: float = 0.0,
+                 kernel_size: int = 9, dilations=(1, 2, 4, 8, 16),
+                 use_global_token: bool = True):
+        super().__init__()
+        h = hidden_dim
+        pad = lambda d: d * (kernel_size // 2)
+
+        self.dw_B = nn.ModuleList([
+            nn.Conv1d(h, h, kernel_size, padding=pad(d), dilation=d, groups=h, bias=False)
+            for d in dilations
+        ])
+        self.mix_B = nn.Conv1d(h * len(dilations), h, 1)
+
+        self.dw_A = nn.ModuleList([
+            nn.Conv1d(h, h, kernel_size, padding=pad(d), dilation=d, groups=h, bias=False)
+            for d in dilations
+        ])
+        self.mix_A = nn.Conv1d(h * len(dilations), h, 1)
+
+        self.proj_B2A = nn.Linear(h, h)
+        self.proj_A2B = nn.Linear(h, h)
+
+        self.use_global_token = use_global_token
+        if use_global_token:
+            self.glb_B2A = nn.Linear(h, h)
+            self.glb_A2B = nn.Linear(h, h)
+
+        self.gate = nn.Linear(h * 4, h * 2)
+        self.dropout = nn.Dropout(dropout)
+        self.normA = nn.LayerNorm(h)
+        self.normB = nn.LayerNorm(h)
+
+    @staticmethod
+    def _agg(x: torch.Tensor, branches: nn.ModuleList, mix: nn.Module) -> torch.Tensor:
+        xch = x.transpose(1, 2)
+        ys = [conv(xch) for conv in branches]
+        y = torch.cat(ys, dim=1)
+        y = mix(y).transpose(1, 2).contiguous()
+        return y
+
+    def forward(self, xA: torch.Tensor, xB: torch.Tensor):
+        ctxB = self._agg(xB, self.dw_B, self.mix_B)
+        ctxA = self._agg(xA, self.dw_A, self.mix_A)
+
+        locA = self.proj_B2A(xB + ctxB)
+        locB = self.proj_A2B(xA + ctxA)
+
+        if self.use_global_token:
+            gB = self.glb_B2A(xB.mean(dim=1, keepdim=True))
+            gA = self.glb_A2B(xA.mean(dim=1, keepdim=True))
+            locA = locA + gB.expand(-1, xA.size(1), -1)
+            locB = locB + gA.expand(-1, xB.size(1), -1)
+
+        z = torch.cat([xA, xB, xA - xB, xA * xB], dim=-1)
+        gA, gB = self.gate(z).chunk(2, dim=-1)
+        gA = torch.sigmoid(gA)
+        gB = torch.sigmoid(gB)
+
+        yA = self.normA(xA + self.dropout(gA * locA))
+        yB = self.normB(xB + self.dropout(gB * locB))
+        return yA, yB
+
+
+def semantic_preservation_loss(R_plus: torch.Tensor, H_S_plus: torch.Tensor,
+                               λ_recon: float = 1.0, λ_local: float = 0.5, λ_global: float = 0.2):
+    recon = F.mse_loss(H_S_plus, R_plus)
+
+    if R_plus.size(1) >= 2:
+        d_ref = R_plus[:, 1:] - R_plus[:, :-1]
+        d_S = H_S_plus[:, 1:] - H_S_plus[:, :-1]
+        local = F.mse_loss(d_S, d_ref)
+    else:
+        local = torch.tensor(0.0, device=R_plus.device)
+
+    def gram_norm(x):
+        G = torch.einsum("b i d, b j d -> b i j", x, x)
+        return G / (G.norm(dim=(1, 2), keepdim=True) + 1e-6)
+
+    glob = F.mse_loss(gram_norm(H_S_plus), gram_norm(R_plus))
+    return λ_recon * recon + λ_local * local + λ_global * glob
+
+
+@contextmanager
+def eval_mode(*modules):
+    states = [m.training for m in modules]
+    try:
+        for m in modules:
+            if m is not None:
+                m.eval()
+        yield
+    finally:
+        for m, s in zip(modules, states):
+            if m is not None:
+                m.train(s)
+
+
+class SSScanDNAHybridModel(nn.Module):
+    """
+    Latest training-structure CrossDNA backbone.
+
+    HF wrapper strategy:
+    - backbone remains compact 5-token A/C/G/T/N model
+    - HF wrapper maps tokenizer ids -> compact ids
+    - HF wrapper expands 5-way logits -> tokenizer vocab logits
+    """
+
+    def __init__(
+        self,
+        config: Optional[Any] = None,
+        alphabet_size=5,
+        d_model=128,
+        block_size=2048,
+        comba_cfg=None,
+        transformer_cfg=None,
+        depth=4,
+        drop_path_rates=None,
+        pretrain=False,
+        for_representation=False,
+        use_final_conv=False,
+        use_s_scan: bool = True,
+        use_mem: bool = False,
+        use_rc_kl: bool = False,
+        use_barlow: bool = False,
+        use_tv: bool = False,
+        sem_max_weight: float = 0.2,
+        sem_warmup_steps: int = 3000,
+        rc_max_weight: float = 0.2,
+        rc_warmup_steps: int = 2000,
+        rc_tau: float = 1.5,
+        rc_bidirectional_stopgrad: bool = True,
+        aux_ce_weight: float = 0.1,
+        gate_freeze_steps: int = 1000,
+        detach_gate: bool = False,
+        gate_sup_weight: float = 0.005,
+        gate_sup_warmup_steps: int = 500,
+        gate_temp: float = 2.0,
+        dropout=0.1,
+        use_ema_teacher: bool = True,
+        ema_decay: float = 0.999,
+        auto_update_ema_in_forward: bool = True,
+        use_bridge: bool = True,
+        bridge_dropout: float = 0.0,
+        use_checkpointing: bool = True,
+        checkpoint_chunk_size: int = 2,
+        checkpoint_core_layers: bool = False,
+        core_checkpoint_chunk_size: int = 1,
+        return_ab_logits: bool = True,
+        streaming_loss: bool = True,
+        streaming_report_ab: bool = True,
+        disable_cross_view: bool = False,
+        **unused_kwargs,
+    ):
+        super().__init__()
+
+        self.config = config
+        if config is not None:
+            cfg = _to_plain_container(config)
+
+            alphabet_size = _cfg_get(cfg, "alphabet_size", alphabet_size)
+            d_model = _cfg_get(cfg, "d_model", d_model)
+            block_size = _cfg_get(cfg, "block_size", block_size)
+            depth = _cfg_get(cfg, "depth", depth)
+            drop_path_rates = _cfg_get(cfg, "drop_path_rates", drop_path_rates)
+
+            pretrain = _cfg_get(cfg, "pretrain", pretrain)
+            for_representation = _cfg_get(cfg, "for_representation", for_representation)
+            use_final_conv = _cfg_get(cfg, "use_final_conv", use_final_conv)
+
+            use_s_scan = _cfg_get(cfg, "use_s_scan", use_s_scan)
+            use_mem = _cfg_get(cfg, "use_mem", use_mem)
+            use_rc_kl = _cfg_get(cfg, "use_rc_kl", use_rc_kl)
+            use_barlow = _cfg_get(cfg, "use_barlow", use_barlow)
+            use_tv = _cfg_get(cfg, "use_tv", use_tv)
+
+            sem_max_weight = _cfg_get(cfg, "sem_max_weight", sem_max_weight)
+            sem_warmup_steps = _cfg_get(cfg, "sem_warmup_steps", sem_warmup_steps)
+            rc_max_weight = _cfg_get(cfg, "rc_max_weight", rc_max_weight)
+            rc_warmup_steps = _cfg_get(cfg, "rc_warmup_steps", rc_warmup_steps)
+            rc_tau = _cfg_get(cfg, "rc_tau", rc_tau)
+            rc_bidirectional_stopgrad = _cfg_get(cfg, "rc_bidirectional_stopgrad", rc_bidirectional_stopgrad)
+
+            aux_ce_weight = _cfg_get(cfg, "aux_ce_weight", aux_ce_weight)
+            gate_freeze_steps = _cfg_get(cfg, "gate_freeze_steps", gate_freeze_steps)
+            detach_gate = _cfg_get(cfg, "detach_gate", detach_gate)
+            gate_sup_weight = _cfg_get(cfg, "gate_sup_weight", gate_sup_weight)
+            gate_sup_warmup_steps = _cfg_get(cfg, "gate_sup_warmup_steps", gate_sup_warmup_steps)
+            gate_temp = _cfg_get(cfg, "gate_temp", gate_temp)
+            dropout = _cfg_get(cfg, "dropout", dropout)
+
+            use_bridge = _cfg_get(cfg, "use_bridge", use_bridge)
+            bridge_dropout = _cfg_get(cfg, "bridge_dropout", bridge_dropout)
+
+            use_checkpointing = _cfg_get(cfg, "use_checkpointing", use_checkpointing)
+            checkpoint_chunk_size = _cfg_get(cfg, "checkpoint_chunk_size", checkpoint_chunk_size)
+
+            checkpoint_core_layers = _cfg_get(cfg, "checkpoint_core_layers", checkpoint_core_layers)
+            core_checkpoint_chunk_size = _cfg_get(cfg, "core_checkpoint_chunk_size", core_checkpoint_chunk_size)
+
+            return_ab_logits = _cfg_get(cfg, "return_ab_logits", return_ab_logits)
+            streaming_loss = _cfg_get(cfg, "streaming_loss", streaming_loss)
+            streaming_report_ab = _cfg_get(cfg, "streaming_report_ab", streaming_report_ab)
+
+            use_ema_teacher = _cfg_get(cfg, "use_ema_teacher", use_ema_teacher)
+            ema_decay = _cfg_get(cfg, "ema_decay", ema_decay)
+            auto_update_ema_in_forward = _cfg_get(cfg, "auto_update_ema_in_forward", auto_update_ema_in_forward)
+            disable_cross_view = _cfg_get(cfg, "disable_cross_view", disable_cross_view)
+
+            transformer_cfg = _cfg_get(cfg, "transformer_cfg", transformer_cfg)
+            comba_cfg = _cfg_get(cfg, "comba_cfg", comba_cfg)
+
+            transformer_cfg = _to_plain_container(transformer_cfg)
+            comba_cfg = _to_plain_container(comba_cfg)
+            drop_path_rates = _to_plain_container(drop_path_rates)
+
+        self.alphabet_size = int(alphabet_size)
+        self.pretrain = bool(pretrain)
+        self.for_representation = bool(for_representation)
+        self.block_size = int(block_size)
+        self.use_final_conv = bool(use_final_conv)
+        self.d_model = int(d_model)
+
+        self.use_checkpointing = bool(use_checkpointing)
+        self.checkpoint_chunk_size = int(checkpoint_chunk_size)
+
+        self.checkpoint_core_layers = bool(checkpoint_core_layers)
+        self.core_checkpoint_chunk_size = int(core_checkpoint_chunk_size)
+
+        self.return_ab_logits = bool(return_ab_logits)
+        self.streaming_loss = bool(streaming_loss)
+        self.streaming_report_ab = bool(streaming_report_ab)
+
+        self.use_ema_teacher = bool(use_ema_teacher)
+        self.ema_decay = float(ema_decay)
+        self.auto_update_ema_in_forward = bool(auto_update_ema_in_forward)
+        self.disable_cross_view = bool(disable_cross_view)
+
+        self.register_buffer("g_step", torch.zeros(1, dtype=torch.long))
+
+        self.linear = nn.Conv1d(self.alphabet_size, self.d_model, kernel_size=9, padding=4)
+        self.rc_linear = nn.Conv1d(self.alphabet_size, self.d_model, kernel_size=9, padding=4)
+
+        self.mlm_mask_embed = nn.Parameter(torch.zeros(self.d_model))
+        self.special_token_embed = nn.Parameter(torch.zeros(self.d_model))
+        nn.init.normal_(self.mlm_mask_embed, mean=0.0, std=0.02)
+        nn.init.normal_(self.special_token_embed, mean=0.0, std=0.02)
+
+        self.branchA_core = DeepEnhancedBranch(
+            hidden_dim=self.d_model,
+            comba_cfg=comba_cfg,
+            transformer_cfg=transformer_cfg,
+            depth=int(depth),
+            drop_path_rates=drop_path_rates,
+            checkpoint_core_layers=self.checkpoint_core_layers,
+            core_checkpoint_chunk_size=self.core_checkpoint_chunk_size,
+        )
+        self.branchB_core = DeepEnhancedBranch(
+            hidden_dim=self.d_model,
+            comba_cfg=comba_cfg,
+            transformer_cfg=transformer_cfg,
+            depth=int(depth),
+            drop_path_rates=drop_path_rates,
+            checkpoint_core_layers=self.checkpoint_core_layers,
+            core_checkpoint_chunk_size=self.core_checkpoint_chunk_size,
+        )
+
+        self.use_bridge = bool(use_bridge)
+        if self.use_bridge:
+            self.bridge = TokenBridge(self.d_model, dropout=float(bridge_dropout))
+        else:
+            self.bridge = None
+
+        if self.use_ema_teacher:
+            self.branchA_core_ema = copy.deepcopy(self.branchA_core)
+            self.branchB_core_ema = copy.deepcopy(self.branchB_core)
+            for p in self.branchA_core_ema.parameters():
+                p.requires_grad_(False)
+            for p in self.branchB_core_ema.parameters():
+                p.requires_grad_(False)
+
+            if self.use_bridge:
+                self.bridge_ema = copy.deepcopy(self.bridge)
+                for p in self.bridge_ema.parameters():
+                    p.requires_grad_(False)
+            else:
+                self.bridge_ema = None
+        else:
+            self.branchA_core_ema = None
+            self.branchB_core_ema = None
+            self.bridge_ema = None
+
+        self.proj_A = Mlp(self.d_model, self.d_model * 2, self.d_model, activation=F.gelu, return_residual=True)
+        self.proj_B = Mlp(self.d_model, self.d_model * 2, self.d_model, activation=F.gelu, return_residual=True)
+        self.gate_fuse = nn.Linear(2 * self.d_model, self.d_model)
+        self.out_linear = nn.Linear(self.d_model, self.alphabet_size)
+        self.dropout = nn.Dropout(float(dropout))
+
+        P_comp, _ = make_complement_perm(self.alphabet_size)
+        self.register_buffer("P_comp", P_comp)
+
+        self.use_s_scan = bool(use_s_scan)
+        self.use_mem = bool(use_mem)
+        self.use_rc_kl = bool(use_rc_kl)
+        self.use_barlow = bool(use_barlow)
+        self.use_tv = bool(use_tv)
+
+        self.sem_max_weight = float(sem_max_weight)
+        self.sem_warmup_steps = int(sem_warmup_steps)
+        self.rc_max_weight = float(rc_max_weight)
+        self.rc_warmup_steps = int(rc_warmup_steps)
+        self.rc_tau = float(rc_tau)
+        self.rc_bidirectional_stopgrad = bool(rc_bidirectional_stopgrad)
+
+        self.aux_ce_weight = float(aux_ce_weight)
+        self.gate_freeze_steps = int(gate_freeze_steps)
+        self.detach_gate = bool(detach_gate)
+        self.gate_sup_weight = float(gate_sup_weight)
+        self.gate_sup_warmup_steps = int(gate_sup_warmup_steps)
+        self.gate_temp = float(gate_temp)
+
+        if self.use_final_conv:
+            self.final_conv = nn.Conv1d(self.d_model, self.d_model, kernel_size=3, padding=1)
+
+        self._unused_init_kwargs = dict(unused_kwargs) if unused_kwargs else {}
+
+    def _branch_receives_rc(self, branch: str, block_idx: int) -> bool:
+        if self.disable_cross_view:
+            return branch == "B"
+
+        if branch == "A":
+            return (block_idx % 2) == 1
+        if branch == "B":
+            return (block_idx % 2) == 0
+        raise ValueError(f"Unknown branch: {branch}")
+
+    def _route_block_inputs(self, t_block: int, fwd_in: torch.Tensor, rc_in: torch.Tensor):
+        if self._branch_receives_rc("A", t_block):
+            return rc_in, fwd_in
+        return fwd_in, rc_in
+
+    def _realign_chunk_outputs(self, H_A: torch.Tensor, H_B: torch.Tensor, chunk_start: int):
+        num_blocks = H_A.size(0)
+        for c in range(num_blocks):
+            actual_t = chunk_start + c
+            if self._branch_receives_rc("A", actual_t):
+                H_A[c] = torch.flip(H_A[c], dims=[1])
+            if self._branch_receives_rc("B", actual_t):
+                H_B[c] = torch.flip(H_B[c], dims=[1])
+        return H_A, H_B
+
+    def _build_chunk_rc_masks(self, chunk_start: int, num_blocks: int, B: int, L_blk: int, device):
+        rcA = torch.tensor(
+            [self._branch_receives_rc("A", chunk_start + c) for c in range(num_blocks)],
+            device=device,
+            dtype=torch.bool,
+        )
+        maskA_row = rcA.repeat_interleave(B).unsqueeze(1)
+        maskA = maskA_row.expand(-1, L_blk)
+        maskB = ~maskA
+        return maskA, maskB
+
+    @torch.no_grad()
+    def update_ema(self):
+        if not self.use_ema_teacher:
+            return
+        if self.branchA_core_ema is None or self.branchB_core_ema is None:
+            return
+
+        d = float(self.ema_decay)
+        for m_ema, m in [
+            (self.branchA_core_ema, self.branchA_core),
+            (self.branchB_core_ema, self.branchB_core),
+        ]:
+            for p_ema, p in zip(m_ema.parameters(), m.parameters()):
+                p_ema.data.lerp_(p.data, 1.0 - d)
+
+        if self.use_bridge and (self.bridge is not None) and (self.bridge_ema is not None):
+            for p_ema, p in zip(self.bridge_ema.parameters(), self.bridge.parameters()):
+                p_ema.data.lerp_(p.data, 1.0 - d)
+
+    def _forward_s_scan_chunk_streaming(
+        self,
+        X_A: torch.Tensor,
+        X_B: torch.Tensor,
+        A_emb_fwd: torch.Tensor,
+        B_emb_rc: torch.Tensor,
+        mlm_labels: torch.Tensor,
+        chunk_start_t: torch.Tensor,
+        num_blocks_t: torch.Tensor,
+        step_t: torch.Tensor,
+        report_ab_t: torch.Tensor,
+    ):
+        chunk_start = int(chunk_start_t.item())
+        num_blocks = int(num_blocks_t.item())
+        step = int(step_t.item())
+        report_ab = bool(int(report_ab_t.item()))
+
+        BC, L_blk, H = X_A.shape
+        B = BC // max(1, num_blocks)
+        device = X_A.device
+
+        H_A = self.branchA_core(X_A)
+        H_B = self.branchB_core(X_B)
+
+        H_A = H_A.view(num_blocks, B, L_blk, H)
+        H_B = H_B.view(num_blocks, B, L_blk, H)
+        H_A, H_B = self._realign_chunk_outputs(H_A, H_B, chunk_start)
+        H_A = H_A.reshape(BC, L_blk, H)
+        H_B = H_B.reshape(BC, L_blk, H)
+
+        if self.use_bridge and self.bridge is not None:
+            H_A, H_B = self.bridge(H_A, H_B)
+
+        fA, rA = self.proj_A(H_A)
+        FA = fA + rA
+        fB, rB = self.proj_B(H_B)
+        FB = fB + rB
+
+        gate_in = torch.cat([FA, FB], dim=-1)
+        g_logits = self.gate_fuse(gate_in)
+        g_raw = torch.sigmoid(g_logits / max(1e-6, self.gate_temp))
+
+        if step < self.gate_freeze_steps:
+            g = 0.5 * torch.ones_like(g_raw)
+        else:
+            g = g_raw
+
+        if self.detach_gate:
+            mix = g.detach() * FA + (1 - g.detach()) * FB
+        else:
+            mix = g * FA + (1 - g) * FB
+
+        fused = F.layer_norm(mix, (mix.size(-1),))
+        fused = ensure_finite(fused, "fused_blk")
+
+        if self.use_final_conv:
+            fused = self.final_conv(fused.permute(0, 2, 1)).permute(0, 2, 1)
+
+        logits = self.out_linear(fused)
+        C = logits.size(-1)
+
+        logits2d = logits.reshape(-1, C)
+        labels1d = mlm_labels.reshape(-1)
+
+        ce_sum = F.cross_entropy(logits2d, labels1d, ignore_index=-100, reduction="sum")
+
+        with torch.no_grad():
+            valid = (labels1d != -100)
+            n_masked = valid.sum()
+
+        with torch.no_grad():
+            correct1 = torch.zeros([], device=device, dtype=torch.long)
+            correct3 = torch.zeros([], device=device, dtype=torch.long)
+            if n_masked.item() > 0:
+                sel_logits = logits2d[valid]
+                sel_labels = labels1d[valid]
+                pred1 = sel_logits.argmax(dim=-1)
+                correct1 = pred1.eq(sel_labels).sum()
+                top3 = sel_logits.topk(3, dim=-1).indices
+                correct3 = top3.eq(sel_labels.unsqueeze(-1)).any(dim=-1).sum()
+
+        total_aux = torch.zeros([], device=device, dtype=torch.float32)
+
+        if self.pretrain:
+            maskA, maskB = self._build_chunk_rc_masks(chunk_start, num_blocks, B, L_blk, device)
+
+            need_sem = self.sem_max_weight > 0.0
+            if need_sem:
+                with torch.no_grad():
+                    teacherA = self.branchA_core_ema if self.use_ema_teacher else self.branchA_core
+                    teacherB = self.branchB_core_ema if self.use_ema_teacher else self.branchB_core
+                    tbridge = self.bridge_ema if (self.use_bridge and self.use_ema_teacher and self.bridge_ema is not None) else (
+                        self.bridge if self.use_bridge else None
+                    )
+
+                    mods = [teacherA, teacherB] + ([tbridge] if tbridge is not None else [])
+                    with eval_mode(*mods):
+                        R_plus_A = teacherA(A_emb_fwd)
+                        R_plus_B = teacherB(A_emb_fwd)
+                        if tbridge is not None:
+                            R_plus_A, R_plus_B = tbridge(R_plus_A, R_plus_B)
+
+                        R_minus_A_rc = teacherA(B_emb_rc)
+                        R_minus_B_rc = teacherB(B_emb_rc)
+                        R_minus_A_fwd = torch.flip(R_minus_A_rc, dims=[1])
+                        R_minus_B_fwd = torch.flip(R_minus_B_rc, dims=[1])
+                        if tbridge is not None:
+                            R_minus_A_fwd, R_minus_B_fwd = tbridge(R_minus_A_fwd, R_minus_B_fwd)
+
+                R_A_teacher = torch.where(maskA.unsqueeze(-1), R_minus_A_fwd, R_plus_A)
+                R_B_teacher = torch.where(maskB.unsqueeze(-1), R_minus_B_fwd, R_plus_B)
+
+                sem_A = semantic_preservation_loss(R_A_teacher.float(), FA.float())
+                sem_B = semantic_preservation_loss(R_B_teacher.float(), FB.float())
+                w_sem = linear_warmup_weight(step, self.sem_warmup_steps, self.sem_max_weight)
+                total_aux = total_aux + w_sem * (sem_A + sem_B)
+
+            if (self.gate_sup_weight > 0.0) and (step >= self.gate_freeze_steps):
+                g_target = (~maskA).float().unsqueeze(-1)
+                g_token_logits = g_logits.mean(dim=-1, keepdim=True) / max(1e-6, self.gate_temp)
+                w_gate = linear_warmup_weight(
+                    step - self.gate_freeze_steps,
+                    self.gate_sup_warmup_steps,
+                    self.gate_sup_weight,
+                )
+                total_aux = total_aux + w_gate * F.binary_cross_entropy_with_logits(
+                    g_token_logits.float(), g_target.float()
+                )
+
+            need_rc = bool(self.use_rc_kl and (self.rc_max_weight > 0.0))
+            need_ab = bool(report_ab or need_rc)
+
+            logitsA = logitsB = None
+            if need_ab:
+                logitsA = self.out_linear(FA)
+                logitsB = self.out_linear(FB)
+
+            if need_rc:
+                if self.rc_bidirectional_stopgrad:
+                    rc = rc_consistency_bidirectional_stopgrad(logitsA, logitsB, self.P_comp, tau=self.rc_tau)
+                else:
+                    rc = rc_consistency_kl(logitsA, logitsB, self.P_comp, tau=self.rc_tau)
+                w_rc = linear_warmup_weight(step, self.rc_warmup_steps, self.rc_max_weight)
+                total_aux = total_aux + w_rc * rc
+
+            if self.use_barlow:
+                total_aux = total_aux + barlow_strand_loss_v2(H_A.float(), H_B.float())
+            if self.use_tv:
+                total_aux = total_aux + tv_mixed(fused.float())
+
+        with torch.no_grad():
+            correctA1 = torch.zeros([], device=device, dtype=torch.long)
+            correctB1 = torch.zeros([], device=device, dtype=torch.long)
+            correctA3 = torch.zeros([], device=device, dtype=torch.long)
+            correctB3 = torch.zeros([], device=device, dtype=torch.long)
+
+            if report_ab and n_masked.item() > 0:
+                if ('logitsA' not in locals()) or (logitsA is None):
+                    logitsA = self.out_linear(FA)
+                    logitsB = self.out_linear(FB)
+
+                _, perm = make_complement_perm(C, device=device)
+
+                valid = (labels1d != -100)
+                labels_safe = labels1d.clamp_min(0)
+                labels_comp = perm[labels_safe]
+
+                maskA_tok, maskB_tok = self._build_chunk_rc_masks(chunk_start, num_blocks, B, L_blk, device)
+                maskA_tok = maskA_tok.reshape(-1)
+                maskB_tok = maskB_tok.reshape(-1)
+
+                yA = torch.where(maskA_tok, labels_comp, labels_safe)[valid]
+                yB = torch.where(maskB_tok, labels_comp, labels_safe)[valid]
+
+                A2d = logitsA.reshape(-1, C)[valid]
+                B2d = logitsB.reshape(-1, C)[valid]
+
+                predA1 = A2d.argmax(dim=-1)
+                predB1 = B2d.argmax(dim=-1)
+                correctA1 = predA1.eq(yA).sum()
+                correctB1 = predB1.eq(yB).sum()
+
+                topA3 = A2d.topk(3, dim=-1).indices
+                topB3 = B2d.topk(3, dim=-1).indices
+                correctA3 = topA3.eq(yA.unsqueeze(-1)).any(dim=-1).sum()
+                correctB3 = topB3.eq(yB.unsqueeze(-1)).any(dim=-1).sum()
+
+        return ce_sum, n_masked, total_aux, correct1, correct3, correctA1, correctB1, correctA3, correctB3
+
+    def _forward_s_scan_chunk(
+        self,
+        X_A: torch.Tensor,
+        X_B: torch.Tensor,
+        A_emb_fwd: torch.Tensor,
+        B_emb_rc: torch.Tensor,
+        chunk_start_t: torch.Tensor,
+        num_blocks_t: torch.Tensor,
+        step_t: torch.Tensor,
+        need_logits_t: torch.Tensor,
+        need_ab_t: torch.Tensor,
+    ):
+        chunk_start = int(chunk_start_t.item())
+        num_blocks = int(num_blocks_t.item())
+        step = int(step_t.item())
+        need_logits = bool(int(need_logits_t.item()))
+        need_ab = bool(int(need_ab_t.item()))
+
+        BC, L_blk, H = X_A.shape
+        B = BC // max(1, num_blocks)
+        device = X_A.device
+
+        H_A = self.branchA_core(X_A)
+        H_B = self.branchB_core(X_B)
+
+        H_A = H_A.view(num_blocks, B, L_blk, H)
+        H_B = H_B.view(num_blocks, B, L_blk, H)
+        H_A, H_B = self._realign_chunk_outputs(H_A, H_B, chunk_start)
+        H_A = H_A.reshape(BC, L_blk, H)
+        H_B = H_B.reshape(BC, L_blk, H)
+
+        if self.use_bridge and self.bridge is not None:
+            H_A, H_B = self.bridge(H_A, H_B)
+
+        fA, rA = self.proj_A(H_A)
+        FA = fA + rA
+        fB, rB = self.proj_B(H_B)
+        FB = fB + rB
+
+        gate_in_blk = torch.cat([FA, FB], dim=-1)
+        g_logits_blk = self.gate_fuse(gate_in_blk)
+        g_raw_blk = torch.sigmoid(g_logits_blk / max(1e-6, self.gate_temp))
+
+        if step < self.gate_freeze_steps:
+            g_blk = 0.5 * torch.ones_like(g_raw_blk)
+        else:
+            g_blk = g_raw_blk
+
+        if self.detach_gate:
+            mix_blk = g_blk.detach() * FA + (1 - g_blk.detach()) * FB
+        else:
+            mix_blk = g_blk * FA + (1 - g_blk) * FB
+
+        fused_blk = F.layer_norm(mix_blk, (mix_blk.size(-1),))
+        fused_blk = ensure_finite(fused_blk, "fused_blk")
+
+        if self.use_final_conv:
+            fused_blk = self.final_conv(fused_blk.permute(0, 2, 1)).permute(0, 2, 1)
+
+        logits_blk = self.out_linear(fused_blk) if need_logits else fused_blk.new_empty((0,))
+
+        need_rc_logits = bool(self.use_rc_kl and (self.rc_max_weight > 0.0))
+        need_ab_internal = bool(need_ab or need_rc_logits)
+
+        logitsA_blk = self.out_linear(FA) if need_ab_internal else fused_blk.new_empty((0,))
+        logitsB_blk = self.out_linear(FB) if need_ab_internal else fused_blk.new_empty((0,))
+
+        total_aux_blk = torch.zeros([], device=device, dtype=torch.float32)
+
+        if self.pretrain:
+            maskA, maskB = self._build_chunk_rc_masks(chunk_start, num_blocks, B, L_blk, device)
+
+            need_sem = self.sem_max_weight > 0.0
+            if need_sem:
+                with torch.no_grad():
+                    teacherA = self.branchA_core_ema if self.use_ema_teacher else self.branchA_core
+                    teacherB = self.branchB_core_ema if self.use_ema_teacher else self.branchB_core
+                    tbridge = self.bridge_ema if (self.use_bridge and self.use_ema_teacher and self.bridge_ema is not None) else (
+                        self.bridge if self.use_bridge else None
+                    )
+
+                    mods = [teacherA, teacherB] + ([tbridge] if tbridge is not None else [])
+                    with eval_mode(*mods):
+                        R_plus_A = teacherA(A_emb_fwd)
+                        R_plus_B = teacherB(A_emb_fwd)
+                        if tbridge is not None:
+                            R_plus_A, R_plus_B = tbridge(R_plus_A, R_plus_B)
+
+                        R_minus_A_rc = teacherA(B_emb_rc)
+                        R_minus_B_rc = teacherB(B_emb_rc)
+                        R_minus_A_fwd = torch.flip(R_minus_A_rc, dims=[1])
+                        R_minus_B_fwd = torch.flip(R_minus_B_rc, dims=[1])
+                        if tbridge is not None:
+                            R_minus_A_fwd, R_minus_B_fwd = tbridge(R_minus_A_fwd, R_minus_B_fwd)
+
+                R_A_teacher = torch.where(maskA.unsqueeze(-1), R_minus_A_fwd, R_plus_A)
+                R_B_teacher = torch.where(maskB.unsqueeze(-1), R_minus_B_fwd, R_plus_B)
+
+                sem_A = semantic_preservation_loss(R_A_teacher.float(), FA.float())
+                sem_B = semantic_preservation_loss(R_B_teacher.float(), FB.float())
+                w_sem = linear_warmup_weight(step, self.sem_warmup_steps, self.sem_max_weight)
+                total_aux_blk = total_aux_blk + w_sem * (sem_A + sem_B)
+
+            if (self.gate_sup_weight > 0.0) and (step >= self.gate_freeze_steps):
+                g_target_blk = (~maskA).float().unsqueeze(-1)
+                g_token_logits_blk = g_logits_blk.mean(dim=-1, keepdim=True) / max(1e-6, self.gate_temp)
+                w_gate = linear_warmup_weight(
+                    step - self.gate_freeze_steps,
+                    self.gate_sup_warmup_steps,
+                    self.gate_sup_weight,
+                )
+                total_aux_blk = total_aux_blk + w_gate * F.binary_cross_entropy_with_logits(
+                    g_token_logits_blk.float(),
+                    g_target_blk.float(),
+                )
+
+            if self.use_rc_kl and (self.rc_max_weight > 0.0):
+                if logitsA_blk.numel() == 0:
+                    logitsA_blk = self.out_linear(FA)
+                if logitsB_blk.numel() == 0:
+                    logitsB_blk = self.out_linear(FB)
+
+                if self.rc_bidirectional_stopgrad:
+                    rc = rc_consistency_bidirectional_stopgrad(logitsA_blk, logitsB_blk, self.P_comp, tau=self.rc_tau)
+                else:
+                    rc = rc_consistency_kl(logitsA_blk, logitsB_blk, self.P_comp, tau=self.rc_tau)
+
+                w_rc = linear_warmup_weight(step, self.rc_warmup_steps, self.rc_max_weight)
+                total_aux_blk = total_aux_blk + w_rc * rc
+
+            if self.use_barlow:
+                total_aux_blk = total_aux_blk + barlow_strand_loss_v2(H_A.float(), H_B.float())
+            if self.use_tv:
+                total_aux_blk = total_aux_blk + tv_mixed(fused_blk.float())
+
+        return fused_blk, logits_blk, logitsA_blk, logitsB_blk, total_aux_blk
+
+    def forward(self, seq, t=None, cls=None, return_embedding=False, state=None, mask=None, **kwargs):
+        step = int(self.g_step.item())
+        if self.training:
+            self.g_step += 1
+
+        _ = mask
+
+        mlm_mask = None
+        mlm_labels = None
+        special_mask = None
+
+        if self.pretrain:
+            if isinstance(seq, (tuple, list)):
+                mlm_mask = seq[1] if len(seq) >= 2 else None
+                mlm_labels = seq[2] if len(seq) >= 3 else None
+                special_mask = seq[3] if len(seq) >= 4 else None
+                seq = seq[0]
+
+        device_type = seq.device.type if seq.device.type in ["cuda", "cpu", "xpu"] else "cuda"
+        amp_dtype = preferred_amp_dtype()
+
+        rc_seq = reverse_complement(seq)
+
+        with torch.autocast(device_type=device_type, dtype=amp_dtype, enabled=(device_type == "cuda")):
+            seq_oh = one_hot_float(seq, self.alphabet_size, dtype=amp_dtype)
+            rc_oh = one_hot_float(rc_seq, self.alphabet_size, dtype=amp_dtype)
+
+            if special_mask is not None:
+                special_mask = special_mask.to(dtype=torch.bool, device=seq.device)
+                rc_special_mask = torch.flip(special_mask, dims=[1])
+
+                seq_oh = seq_oh.masked_fill(special_mask.unsqueeze(-1), 0.0)
+                rc_oh = rc_oh.masked_fill(rc_special_mask.unsqueeze(-1), 0.0)
+            else:
+                rc_special_mask = None
+
+            h = F.gelu(self.linear(seq_oh.permute(0, 2, 1)))
+            rc_h = F.gelu(self.rc_linear(rc_oh.permute(0, 2, 1)))
+            del seq_oh, rc_oh
+
+            if special_mask is not None:
+                non_special = (~special_mask).to(dtype=h.dtype).unsqueeze(1)
+                rc_non_special = (~rc_special_mask).to(dtype=rc_h.dtype).unsqueeze(1)
+                h = h * non_special
+                rc_h = rc_h * rc_non_special
+
+            if mlm_mask is not None:
+                mlm_mask_f = mlm_mask.to(dtype=h.dtype, device=h.device).unsqueeze(1)
+                rc_mlm_mask_f = torch.flip(mlm_mask, dims=[1]).to(dtype=rc_h.dtype, device=rc_h.device).unsqueeze(1)
+                h = h + mlm_mask_f * self.mlm_mask_embed.view(1, -1, 1)
+                rc_h = rc_h + rc_mlm_mask_f * self.mlm_mask_embed.view(1, -1, 1)
+
+            if special_mask is not None:
+                special_mask_f = special_mask.to(dtype=h.dtype, device=h.device).unsqueeze(1)
+                rc_special_mask_f = rc_special_mask.to(dtype=rc_h.dtype, device=rc_h.device).unsqueeze(1)
+                h = h + special_mask_f * self.special_token_embed.view(1, -1, 1)
+                rc_h = rc_h + rc_special_mask_f * self.special_token_embed.view(1, -1, 1)
+
+            use_streaming = bool(
+                self.pretrain
+                and self.use_s_scan
+                and self.streaming_loss
+                and (mlm_labels is not None)
+                and (not self.for_representation)
+            )
+
+            if use_streaming:
+                B, H, L = h.shape
+                l = self.block_size
+                K = (L + l - 1) // l
+                chunk_size = max(1, self.checkpoint_chunk_size)
+
+                ce_sum_total = torch.zeros([], device=seq.device, dtype=torch.float32)
+                n_total = torch.zeros([], device=seq.device, dtype=torch.long)
+                total_aux = torch.zeros([], device=seq.device, dtype=torch.float32)
+
+                correct1 = torch.zeros([], device=seq.device, dtype=torch.long)
+                correct3 = torch.zeros([], device=seq.device, dtype=torch.long)
+                correctA1 = torch.zeros([], device=seq.device, dtype=torch.long)
+                correctB1 = torch.zeros([], device=seq.device, dtype=torch.long)
+                correctA3 = torch.zeros([], device=seq.device, dtype=torch.long)
+                correctB3 = torch.zeros([], device=seq.device, dtype=torch.long)
+
+                keep_rate = mlm_mask.float().mean() if mlm_mask is not None else torch.tensor(1.0, device=seq.device)
+                report_ab_t = torch.tensor(int(self.streaming_report_ab), device=seq.device)
+
+                for chunk_start in range(0, K, chunk_size):
+                    chunk_end = min(chunk_start + chunk_size, K)
+
+                    X_A_batch, X_B_batch = [], []
+                    Aemb_batch, Bemb_batch = [], []
+                    labels_batch = []
+                    lengths = []
+
+                    for t_block in range(chunk_start, chunk_end):
+                        start = t_block * l
+                        end = min(start + l, L)
+                        blk_len = end - start
+                        lengths.append(blk_len)
+
+                        fwd_emb = h[:, :, start:end].transpose(1, 2).contiguous()
+                        rc_emb = rc_h[:, :, start:end].transpose(1, 2).contiguous()
+
+                        fwd_in = self.dropout(h[:, :, start:end]).transpose(1, 2).contiguous()
+                        rc_in = self.dropout(rc_h[:, :, start:end]).transpose(1, 2).contiguous()
+
+                        XA, XB = self._route_block_inputs(t_block, fwd_in, rc_in)
+                        X_A_batch.append(XA)
+                        X_B_batch.append(XB)
+
+                        Aemb_batch.append(fwd_emb)
+                        Bemb_batch.append(rc_emb)
+                        labels_batch.append(mlm_labels[:, start:end])
+
+                    if len(set(lengths)) == 1:
+                        nb = len(X_A_batch)
+
+                        X_A_tensor = torch.cat(X_A_batch, dim=0)
+                        X_B_tensor = torch.cat(X_B_batch, dim=0)
+                        Aemb_tensor = torch.cat(Aemb_batch, dim=0)
+                        Bemb_tensor = torch.cat(Bemb_batch, dim=0)
+                        labels_tensor = torch.cat(labels_batch, dim=0)
+
+                        if self.training and self.use_checkpointing:
+                            ce_sum, n_masked, aux_blk, c1, c3, a1, b1, a3, b3 = cp.checkpoint(
+                                self._forward_s_scan_chunk_streaming,
+                                X_A_tensor, X_B_tensor, Aemb_tensor, Bemb_tensor, labels_tensor,
+                                torch.tensor(chunk_start, device=seq.device),
+                                torch.tensor(nb, device=seq.device),
+                                torch.tensor(step, device=seq.device),
+                                report_ab_t,
+                                use_reentrant=False,
+                            )
+                        else:
+                            ce_sum, n_masked, aux_blk, c1, c3, a1, b1, a3, b3 = self._forward_s_scan_chunk_streaming(
+                                X_A_tensor, X_B_tensor, Aemb_tensor, Bemb_tensor, labels_tensor,
+                                torch.tensor(chunk_start, device=seq.device),
+                                torch.tensor(nb, device=seq.device),
+                                torch.tensor(step, device=seq.device),
+                                report_ab_t,
+                            )
+
+                        ce_sum_total = ce_sum_total + ce_sum
+                        n_total = n_total + n_masked
+                        total_aux = total_aux + aux_blk
+
+                        correct1 += c1
+                        correct3 += c3
+                        correctA1 += a1
+                        correctB1 += b1
+                        correctA3 += a3
+                        correctB3 += b3
+
+                        del X_A_tensor, X_B_tensor, Aemb_tensor, Bemb_tensor, labels_tensor
+
+                    else:
+                        for idx, t_block in enumerate(range(chunk_start, chunk_end)):
+                            if self.training and self.use_checkpointing:
+                                ce_sum, n_masked, aux_blk, c1, c3, a1, b1, a3, b3 = cp.checkpoint(
+                                    self._forward_s_scan_chunk_streaming,
+                                    X_A_batch[idx], X_B_batch[idx], Aemb_batch[idx], Bemb_batch[idx], labels_batch[idx],
+                                    torch.tensor(t_block, device=seq.device),
+                                    torch.tensor(1, device=seq.device),
+                                    torch.tensor(step, device=seq.device),
+                                    report_ab_t,
+                                    use_reentrant=False,
+                                )
+                            else:
+                                ce_sum, n_masked, aux_blk, c1, c3, a1, b1, a3, b3 = self._forward_s_scan_chunk_streaming(
+                                    X_A_batch[idx], X_B_batch[idx], Aemb_batch[idx], Bemb_batch[idx], labels_batch[idx],
+                                    torch.tensor(t_block, device=seq.device),
+                                    torch.tensor(1, device=seq.device),
+                                    torch.tensor(step, device=seq.device),
+                                    report_ab_t,
+                                )
+
+                            ce_sum_total = ce_sum_total + ce_sum
+                            n_total = n_total + n_masked
+                            total_aux = total_aux + aux_blk
+
+                            correct1 += c1
+                            correct3 += c3
+                            correctA1 += a1
+                            correctB1 += b1
+                            correctA3 += a3
+                            correctB3 += b3
+
+                del h, rc_h
+
+                if self.training and self.use_ema_teacher and self.auto_update_ema_in_forward:
+                    self.update_ema()
+
+                HybridOutput = namedtuple("HybridOutput", ["logits"])
+                step_t = torch.tensor(step, device=seq.device, dtype=torch.long)
+
+                stats = torch.stack([
+                    keep_rate.to(torch.float32),
+                    correct1.to(torch.float32),
+                    correct3.to(torch.float32),
+                    correctA1.to(torch.float32),
+                    correctB1.to(torch.float32),
+                    correctA3.to(torch.float32),
+                    correctB3.to(torch.float32),
+                ], dim=0)
+
+                return HybridOutput(logits=(ce_sum_total, n_total, total_aux, stats, step_t)), None
+
+            fused = None
+
+            if self.use_s_scan:
+                B, H, L = h.shape
+                l = self.block_size
+                K = (L + l - 1) // l
+                chunk_size = max(1, self.checkpoint_chunk_size)
+
+                collect_fused = bool(self.for_representation)
+                collect_logits = (not self.for_representation) or self.pretrain
+                need_ab_logits = bool((self.pretrain and self.return_ab_logits) or self.use_rc_kl)
+
+                fused_out = torch.empty((B, L, self.d_model), device=seq.device, dtype=amp_dtype) if collect_fused else None
+                logits_out = torch.empty((B, L, self.alphabet_size), device=seq.device, dtype=amp_dtype) if collect_logits else None
+                logitsA_out = torch.empty((B, L, self.alphabet_size), device=seq.device, dtype=amp_dtype) if need_ab_logits else None
+                logitsB_out = torch.empty((B, L, self.alphabet_size), device=seq.device, dtype=amp_dtype) if need_ab_logits else None
+
+                mask_A_rc = torch.empty((B, L), device=seq.device, dtype=torch.bool)
+                mask_B_rc = torch.empty((B, L), device=seq.device, dtype=torch.bool)
+
+                total_aux = torch.zeros([], device=seq.device, dtype=torch.float32)
+
+                for chunk_start in range(0, K, chunk_size):
+                    chunk_end = min(chunk_start + chunk_size, K)
+
+                    X_A_batch, X_B_batch = [], []
+                    Aemb_batch, Bemb_batch = [], []
+                    lengths = []
+
+                    for t_block in range(chunk_start, chunk_end):
+                        start = t_block * l
+                        end = min(start + l, L)
+                        blk_len = end - start
+                        lengths.append(blk_len)
+
+                        fwd_emb = h[:, :, start:end].transpose(1, 2).contiguous()
+                        rc_emb = rc_h[:, :, start:end].transpose(1, 2).contiguous()
+
+                        fwd_in = self.dropout(h[:, :, start:end]).transpose(1, 2).contiguous()
+                        rc_in = self.dropout(rc_h[:, :, start:end]).transpose(1, 2).contiguous()
+
+                        XA, XB = self._route_block_inputs(t_block, fwd_in, rc_in)
+                        X_A_batch.append(XA)
+                        X_B_batch.append(XB)
+
+                        Aemb_batch.append(fwd_emb)
+                        Bemb_batch.append(rc_emb)
+
+                        mask_A_rc[:, start:end] = self._branch_receives_rc("A", t_block)
+                        mask_B_rc[:, start:end] = self._branch_receives_rc("B", t_block)
+
+                    if len(set(lengths)) == 1:
+                        blk_len = lengths[0]
+                        X_A_tensor = torch.cat(X_A_batch, dim=0)
+                        X_B_tensor = torch.cat(X_B_batch, dim=0)
+                        Aemb_tensor = torch.cat(Aemb_batch, dim=0)
+                        Bemb_tensor = torch.cat(Bemb_batch, dim=0)
+
+                        need_logits_t = torch.tensor(int(collect_logits), device=seq.device)
+                        need_ab_t = torch.tensor(int(need_ab_logits), device=seq.device)
+
+                        if self.training and self.use_checkpointing:
+                            fused_blk, logits_blk, logitsA_blk, logitsB_blk, aux_blk = cp.checkpoint(
+                                self._forward_s_scan_chunk,
+                                X_A_tensor, X_B_tensor, Aemb_tensor, Bemb_tensor,
+                                torch.tensor(chunk_start, device=seq.device),
+                                torch.tensor(len(X_A_batch), device=seq.device),
+                                torch.tensor(step, device=seq.device),
+                                need_logits_t, need_ab_t,
+                                use_reentrant=False,
+                            )
+                        else:
+                            fused_blk, logits_blk, logitsA_blk, logitsB_blk, aux_blk = self._forward_s_scan_chunk(
+                                X_A_tensor, X_B_tensor, Aemb_tensor, Bemb_tensor,
+                                torch.tensor(chunk_start, device=seq.device),
+                                torch.tensor(len(X_A_batch), device=seq.device),
+                                torch.tensor(step, device=seq.device),
+                                need_logits_t, need_ab_t,
+                            )
+
+                        total_aux = total_aux + aux_blk
+
+                        nb = len(X_A_batch)
+                        fused_view = fused_blk.view(nb, B, blk_len, -1)
+                        logits_view = logits_blk.view(nb, B, blk_len, -1) if (collect_logits and logits_blk.numel() > 0) else None
+                        logitsA_view = logitsA_blk.view(nb, B, blk_len, -1) if (need_ab_logits and logitsA_blk.numel() > 0) else None
+                        logitsB_view = logitsB_blk.view(nb, B, blk_len, -1) if (need_ab_logits and logitsB_blk.numel() > 0) else None
+
+                        for c, t_block in enumerate(range(chunk_start, chunk_end)):
+                            start = t_block * l
+                            end = min(start + l, L)
+
+                            if collect_fused:
+                                fused_out[:, start:end, :] = fused_view[c]
+                            if collect_logits:
+                                logits_out[:, start:end, :] = logits_view[c]
+                            if need_ab_logits:
+                                logitsA_out[:, start:end, :] = logitsA_view[c]
+                                logitsB_out[:, start:end, :] = logitsB_view[c]
+
+                        del X_A_tensor, X_B_tensor, Aemb_tensor, Bemb_tensor
+                        del fused_blk, logits_blk, logitsA_blk, logitsB_blk
+
+                    else:
+                        for idx, t_block in enumerate(range(chunk_start, chunk_end)):
+                            start = t_block * l
+                            end = min(start + l, L)
+
+                            X_A_tensor = X_A_batch[idx]
+                            X_B_tensor = X_B_batch[idx]
+                            Aemb_tensor = Aemb_batch[idx]
+                            Bemb_tensor = Bemb_batch[idx]
+
+                            need_logits_t = torch.tensor(int(collect_logits), device=seq.device)
+                            need_ab_t = torch.tensor(int(need_ab_logits), device=seq.device)
+
+                            if self.training and self.use_checkpointing:
+                                fused_blk, logits_blk, logitsA_blk, logitsB_blk, aux_blk = cp.checkpoint(
+                                    self._forward_s_scan_chunk,
+                                    X_A_tensor, X_B_tensor, Aemb_tensor, Bemb_tensor,
+                                    torch.tensor(t_block, device=seq.device),
+                                    torch.tensor(1, device=seq.device),
+                                    torch.tensor(step, device=seq.device),
+                                    need_logits_t, need_ab_t,
+                                    use_reentrant=False,
+                                )
+                            else:
+                                fused_blk, logits_blk, logitsA_blk, logitsB_blk, aux_blk = self._forward_s_scan_chunk(
+                                    X_A_tensor, X_B_tensor, Aemb_tensor, Bemb_tensor,
+                                    torch.tensor(t_block, device=seq.device),
+                                    torch.tensor(1, device=seq.device),
+                                    torch.tensor(step, device=seq.device),
+                                    need_logits_t, need_ab_t,
+                                )
+
+                            total_aux = total_aux + aux_blk
+
+                            if collect_fused:
+                                fused_out[:, start:end, :] = fused_blk
+                            if collect_logits:
+                                logits_out[:, start:end, :] = logits_blk
+                            if need_ab_logits and logitsA_blk.numel() > 0:
+                                logitsA_out[:, start:end, :] = logitsA_blk
+                                logitsB_out[:, start:end, :] = logitsB_blk
+
+                            del fused_blk, logits_blk, logitsA_blk, logitsB_blk
+
+                del h, rc_h
+
+                logits = logits_out if collect_logits else None
+                logits_A_only = logitsA_out if need_ab_logits else None
+                logits_B_only = logitsB_out if need_ab_logits else None
+                fused = fused_out if collect_fused else None
+
+            else:
+                feat = self.dropout(h).transpose(1, 2).contiguous()
+                rc_feat = self.dropout(rc_h).transpose(1, 2).contiguous()
+
+                H_A = self.branchA_core(feat)
+                H_Br = self.branchB_core(rc_feat)
+                R_A = H_A
+                R_B = torch.flip(H_Br, dims=[1])
+
+                if self.use_bridge and self.bridge is not None:
+                    R_A, R_B = self.bridge(R_A, R_B)
+
+                fA, rA = self.proj_A(R_A)
+                FA = fA + rA
+                fB, rB = self.proj_B(R_B)
+                FB = fB + rB
+
+                gate_in = torch.cat([FA, FB], dim=-1)
+                g_logits = self.gate_fuse(gate_in)
+                g_raw = torch.sigmoid(g_logits / max(1e-6, self.gate_temp))
+
+                if step < self.gate_freeze_steps:
+                    g = 0.5 * torch.ones_like(g_raw)
+                else:
+                    g = g_raw
+
+                if self.detach_gate:
+                    mix = g.detach() * FA + (1 - g.detach()) * FB
+                else:
+                    mix = g * FA + (1 - g) * FB
+
+                fused = F.layer_norm(mix, (mix.size(-1),))
+                fused = ensure_finite(fused, "fused")
+
+                if self.use_final_conv:
+                    fused = self.final_conv(fused.permute(0, 2, 1)).permute(0, 2, 1)
+
+                logits = self.out_linear(fused) if (not self.for_representation or self.pretrain) else None
+
+                need_ab_logits = bool((self.pretrain and self.return_ab_logits) or self.use_rc_kl)
+                logits_A_only = self.out_linear(FA) if need_ab_logits else None
+                logits_B_only = self.out_linear(FB) if need_ab_logits else None
+
+                mask_A_rc = torch.zeros(FA.size()[:2], dtype=torch.bool, device=FA.device)
+                mask_B_rc = torch.ones_like(mask_A_rc)
+
+                total_aux = logits.new_zeros(()) if self.pretrain and logits is not None else None
+
+                del h, rc_h, feat, rc_feat
+
+        if self.for_representation:
+            return fused, None
+
+        if self.training and self.use_ema_teacher and self.auto_update_ema_in_forward:
+            self.update_ema()
+
+        if self.pretrain:
+            HybridOutput = namedtuple("HybridOutput", ["logits"])
+            return HybridOutput(
+                logits=(
+                    logits,
+                    mlm_mask,
+                    total_aux,
+                    logits_A_only.detach() if logits_A_only is not None else None,
+                    logits_B_only.detach() if logits_B_only is not None else None,
+                    mask_A_rc.detach() if mask_A_rc is not None else None,
+                    mask_B_rc.detach() if mask_B_rc is not None else None,
+                    int(step),
+                )
+            ), None
+
+        return logits, None
+
+    @property
+    def d_output(self):
+        if getattr(self, "d_model", None) is None:
+            raise NotImplementedError("SequenceModule instantiation must set d_output")
+        return self.d_model
+
+
+class CrossDNAForMaskedLM(PreTrainedModel):
+    config_class = CrossDNAConfig
+    base_model_prefix = "backbone"
+    main_input_name = "input_ids"
+
+    # allow older HF checkpoints to load with warnings rather than hard-failing
+    _keys_to_ignore_on_load_missing = [
+        r"backbone\.mlm_mask_embed",
+        r"backbone\.special_token_embed",
+        r"backbone\.branchA_core_ema\..*",
+        r"backbone\.branchB_core_ema\..*",
+        r"backbone\.bridge_ema\..*",
+    ]
+
+    def __init__(self, config: CrossDNAConfig):
+        super().__init__(config)
+        self.config = config
+        self.backbone = SSScanDNAHybridModel(config=config)
+        self.post_init()
+
+    @property
+    def tokenizer_vocab_size(self) -> int:
+        return int(getattr(self.config, "vocab_size", self.config.alphabet_size))
+
+    @property
+    def dna_token_ids(self) -> Dict[str, int]:
+        cfg_ids = getattr(self.config, "dna_token_ids", None)
+        if cfg_ids is not None:
+            return dict(cfg_ids)
+        start = int(getattr(self.config, "dna_token_start_id", 7))
+        return {"A": start + 0, "C": start + 1, "G": start + 2, "T": start + 3, "N": start + 4}
+
+    @property
+    def compact_n_token_id(self) -> int:
+        return int(getattr(self.config, "compact_n_token_id", self.config.alphabet_size - 1))
+
+    def _to_compact_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        input_ids = input_ids.long()
+        if input_ids.numel() == 0:
+            return input_ids
+
+        mn, mx = int(input_ids.min()), int(input_ids.max())
+        if 0 <= mn and mx < self.config.alphabet_size:
+            return input_ids
+
+        compact = torch.full_like(input_ids, self.compact_n_token_id)
+        for compact_id, base in enumerate(["A", "C", "G", "T", "N"]):
+            tok_id = self.dna_token_ids[base]
+            compact[input_ids == tok_id] = compact_id
+        return compact
+
+    def _labels_to_compact(self, labels: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+        if labels is None:
+            return None
+
+        labels = labels.long()
+        compact = torch.full_like(labels, -100)
+
+        direct_mask = (labels >= 0) & (labels < self.config.alphabet_size)
+        compact[direct_mask] = labels[direct_mask]
+
+        for compact_id, base in enumerate(["A", "C", "G", "T", "N"]):
+            tok_id = self.dna_token_ids[base]
+            compact[labels == tok_id] = compact_id
+
+        return compact
+
+    def _build_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
+        input_ids = input_ids.long()
+        if input_ids.numel() == 0:
+            return torch.zeros_like(input_ids, dtype=torch.bool)
+
+        mn, mx = int(input_ids.min()), int(input_ids.max())
+        if 0 <= mn and mx < self.config.alphabet_size:
+            return torch.zeros_like(input_ids, dtype=torch.bool)
+
+        special_mask = torch.ones_like(input_ids, dtype=torch.bool)
+        for base in ["A", "C", "G", "T", "N"]:
+            special_mask[input_ids == self.dna_token_ids[base]] = False
+        return special_mask
+
+    def _expand_logits_to_tokenizer_vocab(self, compact_logits: torch.Tensor) -> torch.Tensor:
+        B, L, C = compact_logits.shape
+        V = self.tokenizer_vocab_size
+        full_logits = compact_logits.new_full((B, L, V), -1e4)
+
+        for compact_id, base in enumerate(["A", "C", "G", "T", "N"]):
+            full_logits[:, :, self.dna_token_ids[base]] = compact_logits[:, :, compact_id]
+
+        return full_logits
+
+    @torch.no_grad()
+    def extract_embeddings(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        compact_ids = self._to_compact_ids(input_ids)
+
+        was_pretrain = bool(getattr(self.backbone, "pretrain", False))
+        was_for_repr = bool(getattr(self.backbone, "for_representation", False))
+        try:
+            self.backbone.pretrain = False
+            self.backbone.for_representation = True
+            embeddings, _ = self.backbone(compact_ids, mask=attention_mask)
+        finally:
+            self.backbone.pretrain = was_pretrain
+            self.backbone.for_representation = was_for_repr
+        return embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs
+    ):
+        return_dict = self.config.use_return_dict if return_dict is None else return_dict
+
+        compact_input_ids = self._to_compact_ids(input_ids)
+        compact_labels = self._labels_to_compact(labels)
+        special_mask = self._build_special_mask(input_ids)
+
+        if compact_labels is not None:
+            mlm_mask = compact_labels.ne(-100)
+        else:
+            mlm_mask = input_ids.eq(getattr(self.config, "mask_token_id", 3))
+
+        # HF wrapper always asks for dense token logits, so temporarily disable streaming-loss forward.
+        was_streaming = bool(getattr(self.backbone, "streaming_loss", False))
+        self.backbone.streaming_loss = False
+        try:
+            if self.config.pretrain:
+                outputs, _ = self.backbone(
+                    (compact_input_ids, mlm_mask, compact_labels, special_mask),
+                    mask=attention_mask,
+                )
+                compact_logits = outputs.logits[0]
+                aux_loss = outputs.logits[2]
+            else:
+                compact_logits, _ = self.backbone(compact_input_ids, mask=attention_mask)
+                aux_loss = None
+        finally:
+            self.backbone.streaming_loss = was_streaming
+
+        logits = self._expand_logits_to_tokenizer_vocab(compact_logits)
+
+        loss = None
+        if compact_labels is not None:
+            loss = F.cross_entropy(
+                compact_logits.reshape(-1, self.config.alphabet_size),
+                compact_labels.reshape(-1),
+                ignore_index=-100,
+            )
+            if aux_loss is not None:
+                loss = loss + aux_loss.to(loss.dtype)
+
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=None,
+            attentions=None,
+        )

28.6M/huggingface_crossdna_140K_len/crossdna/special_tokens_map.json

@@ -0,0 +1,12 @@
+{
+  "bos_token": "[BOS]",
+  "cls_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]",
+  "additional_special_tokens": [
+    "[RESERVED]"
+  ]
+}

28.6M/huggingface_crossdna_140K_len/crossdna/tokenization_crossdna.py

@@ -0,0 +1,181 @@
+
+import json
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Union
+
+from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
+
+
+class CrossDNATokenizer(PreTrainedTokenizer):
+    def __init__(
+        self,
+        characters: Sequence[str] = ("A", "C", "G", "T", "N"),
+        model_max_length: int = 143360,
+        padding_side: str = "left",
+        dna_token_start_id: int = 7,
+        **kwargs,
+    ):
+        self.characters = [str(ch).upper() for ch in characters]
+        self.model_max_length = int(model_max_length)
+        self.dna_token_start_id = int(dna_token_start_id)
+
+        self._vocab_str_to_int = {
+            "[CLS]": 0,
+            "[SEP]": 1,
+            "[BOS]": 2,
+            "[MASK]": 3,
+            "[PAD]": 4,
+            "[RESERVED]": 5,
+            "[UNK]": 6,
+            **{ch: self.dna_token_start_id + i for i, ch in enumerate(self.characters)},
+        }
+        self._vocab_int_to_str = {v: k for k, v in self._vocab_str_to_int.items()}
+
+        bos_token = AddedToken("[BOS]", lstrip=False, rstrip=False)
+        eos_token = AddedToken("[SEP]", lstrip=False, rstrip=False)
+        sep_token = AddedToken("[SEP]", lstrip=False, rstrip=False)
+        cls_token = AddedToken("[CLS]", lstrip=False, rstrip=False)
+        pad_token = AddedToken("[PAD]", lstrip=False, rstrip=False)
+        unk_token = AddedToken("[UNK]", lstrip=False, rstrip=False)
+        mask_token = AddedToken("[MASK]", lstrip=False, rstrip=False)
+
+        if "add_special_tokens" in kwargs:
+            kwargs.pop("add_special_tokens")
+
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            sep_token=sep_token,
+            cls_token=cls_token,
+            pad_token=pad_token,
+            mask_token=mask_token,
+            unk_token=unk_token,
+            add_prefix_space=False,
+            model_max_length=self.model_max_length,
+            padding_side=padding_side,
+            **kwargs,
+        )
+
+    def __len__(self):
+        return len(self._vocab_str_to_int)
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self._vocab_str_to_int)
+
+    def get_vocab(self) -> Dict[str, int]:
+        return dict(self._vocab_str_to_int)
+
+    def _tokenize(self, text: str) -> List[str]:
+        return list(text.upper())
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self._vocab_str_to_int.get(token, self._vocab_str_to_int["[UNK]"])
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self._vocab_int_to_str.get(index, "[UNK]")
+
+    def convert_tokens_to_string(self, tokens):
+        return "".join(tokens)
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        sep = [self.sep_token_id]
+        result = token_ids_0 + sep
+        if token_ids_1 is not None:
+            result += token_ids_1 + sep
+        return result
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+        already_has_special_tokens: bool = False,
+    ) -> List[int]:
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0,
+                token_ids_1=token_ids_1,
+                already_has_special_tokens=True,
+            )
+        result = ([0] * len(token_ids_0)) + [1]
+        if token_ids_1 is not None:
+            result += ([0] * len(token_ids_1)) + [1]
+        return result
+
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        sep = [self.sep_token_id]
+        result = len(token_ids_0 + sep) * [0]
+        if token_ids_1 is not None:
+            result += len(token_ids_1 + sep) * [1]
+        return result
+
+    def get_config(self) -> Dict:
+        return {
+            "characters": self.characters,
+            "model_max_length": self.model_max_length,
+            "padding_side": self.padding_side,
+            "dna_token_start_id": self.dna_token_start_id,
+            "bos_token": "[BOS]",
+            "eos_token": "[SEP]",
+            "sep_token": "[SEP]",
+            "cls_token": "[CLS]",
+            "pad_token": "[PAD]",
+            "mask_token": "[MASK]",
+            "unk_token": "[UNK]",
+            "tokenizer_class": "CrossDNATokenizer",
+            "auto_map": {
+                "AutoTokenizer": [
+                    "tokenization_crossdna.CrossDNATokenizer",
+                    None
+                ]
+            }
+        }
+
+    @classmethod
+    def from_config(cls, config: Dict) -> "CrossDNATokenizer":
+        return cls(
+            characters=config.get("characters", ["A", "C", "G", "T", "N"]),
+            model_max_length=config.get("model_max_length", 143360),
+            padding_side=config.get("padding_side", "left"),
+            dna_token_start_id=config.get("dna_token_start_id", 7),
+        )
+
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
+        save_directory = Path(save_directory)
+        save_directory.mkdir(parents=True, exist_ok=True)
+
+        cfg_file = save_directory / "tokenizer_config.json"
+        stm_file = save_directory / "special_tokens_map.json"
+
+        with open(cfg_file, "w", encoding="utf-8") as f:
+            json.dump(self.get_config(), f, indent=2, ensure_ascii=False)
+
+        special_tokens_map = {
+            "bos_token": "[BOS]",
+            "cls_token": "[CLS]",
+            "eos_token": "[SEP]",
+            "mask_token": "[MASK]",
+            "pad_token": "[PAD]",
+            "sep_token": "[SEP]",
+            "unk_token": "[UNK]",
+            "additional_special_tokens": ["[RESERVED]"]
+        }
+        with open(stm_file, "w", encoding="utf-8") as f:
+            json.dump(special_tokens_map, f, indent=2, ensure_ascii=False)
+
+        return (str(cfg_file), str(stm_file))
+
+    @classmethod
+    def from_pretrained(cls, save_directory: Union[str, os.PathLike], **kwargs):
+        cfg_file = Path(save_directory) / "tokenizer_config.json"
+        with open(cfg_file, encoding="utf-8") as f:
+            cfg = json.load(f)
+        cfg.update(kwargs)
+        cfg.pop("tokenizer_class", None)
+        cfg.pop("auto_map", None)
+        return cls.from_config(cfg)

28.6M/huggingface_crossdna_140K_len/crossdna/tokenizer_config.json

@@ -0,0 +1,26 @@
+{
+  "characters": [
+    "A",
+    "C",
+    "G",
+    "T",
+    "N"
+  ],
+  "model_max_length": 143360,
+  "padding_side": "left",
+  "dna_token_start_id": 7,
+  "bos_token": "[BOS]",
+  "eos_token": "[SEP]",
+  "sep_token": "[SEP]",
+  "cls_token": "[CLS]",
+  "pad_token": "[PAD]",
+  "mask_token": "[MASK]",
+  "unk_token": "[UNK]",
+  "tokenizer_class": "CrossDNATokenizer",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_crossdna.CrossDNATokenizer",
+      null
+    ]
+  }
+}

28.6M/huggingface_crossdna_140K_len/crossdna/transfer.py

@@ -0,0 +1,81 @@
+# transfer.py
+import os
+import json
+
+# 1) 先关掉 dynamo / compile，必须放在导入 modeling_crossdna 之前
+os.environ["DISABLE_TORCH_COMPILE"] = "1"
+os.environ["TORCHDYNAMO_DISABLE"] = "1"
+
+import torch
+
+# 2) 再做一次硬 monkey-patch，避免 fla 在 import 时触发 @torch.compile
+if hasattr(torch, "compile"):
+    def _no_compile(fn=None, *args, **kwargs):
+        if fn is None:
+            def deco(f):
+                return f
+            return deco
+        return fn
+    torch.compile = _no_compile
+
+from configuration_crossdna import CrossDNAConfig
+from modeling_crossdna import CrossDNAForMaskedLM
+
+CKPT = "/data/zhaol/projects/huggingface_crossdna_140K_len/crossdna/last.ckpt"
+MODEL_DIR = "/data/zhaol/projects/huggingface_crossdna_140K_len/crossdna"
+
+
+def adapt_state_dict(sd):
+    new_sd = {}
+    for k, v in sd.items():
+        k2 = k
+
+        # Lightning / DDP / compiled 常见前缀
+        for pref in (
+            "state_dict.",
+            "model.",
+            "module.",
+            "_orig_mod.",
+        ):
+            if k2.startswith(pref):
+                k2 = k2[len(pref):]
+
+        # 如果 ckpt 里不是 HF 包装后的 key，就挂到 backbone 下
+        if not k2.startswith("backbone."):
+            k2 = "backbone." + k2
+
+        new_sd[k2] = v
+    return new_sd
+
+
+# 1) 读 config.json 构建 HF 模型骨架
+with open(os.path.join(MODEL_DIR, "config.json"), "r", encoding="utf-8") as f:
+    cfg_dict = json.load(f)
+
+cfg = CrossDNAConfig(**cfg_dict)
+model = CrossDNAForMaskedLM(cfg)
+
+# 2) 读 ckpt
+raw = torch.load(CKPT, map_location="cpu")
+sd = raw.get("state_dict", raw) if isinstance(raw, dict) else raw
+sd = adapt_state_dict(sd)
+
+# 3) 加载权重
+missing, unexpected = model.load_state_dict(sd, strict=False)
+
+print("[Missing]", len(missing), "keys")
+if missing:
+    print("  first 30 missing:")
+    for k in missing[:30]:
+        print("   ", k)
+
+print("[Unexpected]", len(unexpected), "keys")
+if unexpected:
+    print("  first 30 unexpected:")
+    for k in unexpected[:30]:
+        print("   ", k)
+
+# 4) 直接保存为 HF safetensors
+model.save_pretrained(MODEL_DIR, safe_serialization=True)
+
+print("Saved HF weights to:", os.path.join(MODEL_DIR, "model.safetensors"))

28.6M/huggingface_crossdna_140K_len/crossdna_140K_infer.py

@@ -0,0 +1,48 @@
+import os
+
+# 必须放在最前面
+os.environ["DISABLE_TORCH_COMPILE"] = "1"
+os.environ["TORCHDYNAMO_DISABLE"] = "1"
+
+import torch
+
+# 必须在 from_pretrained 之前，把 torch.compile 变成 no-op
+if hasattr(torch, "compile"):
+    def _no_compile(fn=None, *args, **kwargs):
+        if fn is None:
+            def deco(f):
+                return f
+            return deco
+        return fn
+    torch.compile = _no_compile
+
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+MODEL_DIR = "/data/zhaol/projects/huggingface_crossdna_140K_len/crossdna"
+
+tok = AutoTokenizer.from_pretrained(
+    MODEL_DIR,
+    trust_remote_code=True,
+    local_files_only=True,
+)
+
+model = AutoModelForMaskedLM.from_pretrained(
+    MODEL_DIR,
+    trust_remote_code=True,
+    local_files_only=True,
+).eval()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+seq = "ACGT" * 128
+enc = tok(seq, return_tensors="pt", add_special_tokens=False)
+x = enc["input_ids"].to(device)
+
+with torch.inference_mode():
+    out = model(input_ids=x)
+    emb = model.extract_embeddings(x)
+
+print("input_ids.shape =", tuple(x.shape))
+print("logits.shape =", tuple(out.logits.shape))
+print("embeddings.shape =", tuple(emb.shape))