src/dual_stream_adapter/adapter.py

"""
dual_stream_adapter/adapter.py

Dual-Stream Adapter with TWO separate pre-trained models.

Content Stream:  DeepSeek-Coder 6.7B (understands code, long sequences)
Context Stream:  Llama 3.2 3B (understands instructions, intentions, plans)

Both models are 4-bit quantized and frozen. Only the projection +
cross-attention gate are trainable.

Usage:
    adapter = DualStreamAdapter(
        content_model="deepseek-ai/deepseek-coder-6.7b-instruct",
        context_model="meta-llama/Llama-3.2-3B-Instruct",
    )
    adapter.load_models()     # 4-bit both
    adapter.freeze_all()      # Only projection + cross + gate train
"""

from __future__ import annotations

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig


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

    def forward(self, x: Tensor) -> Tensor:
        d = x.float()
        return (d * torch.rsqrt(d.pow(2).mean(-1, keepdim=True) + self.eps)).to(x.dtype) * self.weight


class CrossAttentionGate(nn.Module):
    """Content queries context via cross-attention, gated by content-driven sigmoid.

    Projects K/V from context_dim to inner_dim, Q from content_dim to inner_dim.
    Gate uses a bottleneck: content_dim → 256 → content_dim → sigmoid.
    """

    def __init__(self, content_d: int, context_d: int,
                 n_head: int, n_kv: int, head_dim: int):
        super().__init__()
        self.content_d = content_d
        self.n_head, self.n_kv, self.hd = n_head, n_kv, head_dim
        self.d_inner = n_kv * head_dim
        self.groups = n_head // n_kv if n_head >= n_kv else 1
        self.scale = head_dim ** -0.5

        # Q from content, K/V from context — direct to inner_dim
        self.q_proj = nn.Linear(content_d, n_head * head_dim, bias=False)
        self.k_proj = nn.Linear(context_d, self.d_inner, bias=False)
        self.v_proj = nn.Linear(context_d, self.d_inner, bias=False)
        self.o_proj = nn.Linear(n_head * head_dim, content_d, bias=False)

        # Bottleneck gate: content_d → 256 → content_d → sigmoid
        gate_bottleneck = 256
        self.gate_down = nn.Linear(content_d, gate_bottleneck, bias=False)
        self.gate_up = nn.Linear(gate_bottleneck, content_d, bias=False)
        self.norm = RMSNorm(content_d)

        # KV cache (computed from context, reused across steps)
        self._ctx_k: Optional[Tensor] = None
        self._ctx_v: Optional[Tensor] = None
        self._ctx_pad: Optional[Tensor] = None

    def _gate(self, content: Tensor) -> Tensor:
        return torch.sigmoid(self.gate_up(torch.nn.functional.gelu(self.gate_down(content))))

    def cache_context(self, context: Tensor, ctx_pad: Optional[Tensor] = None):
        """Precompute and cache context K/V for incremental generation."""
        B, Tctx, _ = context.shape
        self._ctx_k = self.k_proj(context).view(B, Tctx, self.n_kv, self.hd).transpose(1, 2)
        self._ctx_v = self.v_proj(context).view(B, Tctx, self.n_kv, self.hd).transpose(1, 2)
        self._ctx_k = self._ctx_k.repeat_interleave(self.groups, dim=1)
        self._ctx_v = self._ctx_v.repeat_interleave(self.groups, dim=1)
        if ctx_pad is not None:
            self._ctx_pad = ctx_pad
        else:
            self._ctx_pad = None

    def clear_cache(self):
        self._ctx_k = None
        self._ctx_v = None
        self._ctx_pad = None

    def forward_step(self, content: Tensor, record_gate: bool = False) -> Tensor:
        B, Tc, _ = content.shape
        if self._ctx_k is None:
            raise RuntimeError("Context KV not cached. Call cache_context() first.")
        Tctx = self._ctx_k.shape[2]

        q = self.q_proj(content).view(B, Tc, self.n_head, self.hd).transpose(1, 2)

        mask = self._ctx_pad[:, None, None, :].expand(B, 1, Tc, Tctx) if self._ctx_pad is not None else None
        retrieved = F.scaled_dot_product_attention(q, self._ctx_k, self._ctx_v, attn_mask=mask, scale=self.scale)
        retrieved = self.o_proj(retrieved.transpose(1, 2).contiguous().view(B, Tc, -1))

        gate = self._gate(content)
        if record_gate:
            self._last_gate = gate.detach()
        return self.norm(content + gate * retrieved)

    def forward(self, content: Tensor, context: Tensor,
                ctx_pad: Optional[Tensor] = None) -> Tensor:
        B, Tc, _ = content.shape
        Tctx = context.shape[1]

        q = self.q_proj(content).view(B, Tc, self.n_head, self.hd).transpose(1, 2)
        k = self.k_proj(context).view(B, Tctx, self.n_kv, self.hd).transpose(1, 2)
        v = self.v_proj(context).view(B, Tctx, self.n_kv, self.hd).transpose(1, 2)
        k = k.repeat_interleave(self.groups, dim=1)
        v = v.repeat_interleave(self.groups, dim=1)

        mask = ctx_pad[:, None, None, :].expand(B, 1, Tc, Tctx) if ctx_pad is not None else None
        retrieved = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=self.scale)
        retrieved = self.o_proj(retrieved.transpose(1, 2).contiguous().view(B, Tc, -1))

        gate = self._gate(content)
        return self.norm(content + gate * retrieved)


class DualStreamAdapter(nn.Module):
    """
    Two-model Dual-Stream Adapter.

    Content model — specialized for the task domain (code, data, tool outputs)
    Context model — specialized for instructions, planning, constraints
    """

    def __init__(self, content_model: str, context_model: str,
                 max_context_len: int = 1024, max_memory: Optional[dict] = None):
        super().__init__()
        self.content_model_id = content_model
        self.context_model_id = context_model
        self.max_context_len = max_context_len
        self.max_memory = max_memory

        # Loaded state
        self.content_model: Optional[nn.Module] = None
        self.context_model: Optional[nn.Module] = None
        self.content_tokenizer = None
        self.context_tokenizer = None
        self.content_d: int = 0
        self.context_d: int = 0
        self.n_head: int = 0
        self.n_kv: int = 0
        self.head_dim: int = 0
        self.vocab_size: int = 0
        self.cross_gate: Optional[CrossAttentionGate] = None

    def load_models(self):
        """Load both models in 4-bit quantization."""

        bnb = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
            llm_int8_enable_fp32_cpu_offload=True,
        )

        # ── Content model ─────────────────────────────────────────────
        print(f"Content: {self.content_model_id}")
        self.content_tokenizer = AutoTokenizer.from_pretrained(
            self.content_model_id, trust_remote_code=True)
        if self.content_tokenizer.pad_token is None:
            self.content_tokenizer.pad_token = self.content_tokenizer.eos_token

        self.content_model = AutoModelForCausalLM.from_pretrained(
            self.content_model_id,
            quantization_config=bnb,
            device_map="auto",
            max_memory=self.max_memory,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
        )
        cfg = self.content_model.config
        self.content_d = cfg.hidden_size
        self.n_head = cfg.num_attention_heads
        self.n_kv = getattr(cfg, 'num_key_value_heads', self.n_head)
        self.head_dim = getattr(cfg, 'head_dim', self.content_d // self.n_head)
        self.vocab_size = cfg.vocab_size
        print(f"  d={self.content_d} vocab={self.vocab_size}")

        # ── Context model ─────────────────────────────────────────────
        print(f"Context: {self.context_model_id}")
        self.context_tokenizer = AutoTokenizer.from_pretrained(
            self.context_model_id, trust_remote_code=True)
        if self.context_tokenizer.pad_token is None:
            self.context_tokenizer.pad_token = self.context_tokenizer.eos_token

        self.context_model = AutoModelForCausalLM.from_pretrained(
            self.context_model_id,
            quantization_config=bnb,
            device_map="auto",
            max_memory=self.max_memory,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
        )
        self.context_d = self.context_model.config.hidden_size
        print(f"  d={self.context_d} (content d={self.content_d})")

        # ── Cross-attention gate ──────────────────────────────────────
        self.cross_gate = CrossAttentionGate(
            content_d=self.content_d,
            context_d=self.context_d,
            n_head=self.n_head,
            n_kv=self.n_kv,
            head_dim=64,
        )
        device = next(self.content_model.parameters()).device
        dtype = next(self.content_model.parameters()).dtype
        self.cross_gate = self.cross_gate.to(device=device, dtype=dtype)

    def freeze_all(self):
        """Freeze both base models. Only cross-attention gate trains."""
        for model in [self.content_model, self.context_model]:
            if model is not None:
                model.eval()
                for p in model.parameters():
                    p.requires_grad = False

    def _content_forward(self, content_ids: Tensor,
                         attention_mask: Optional[Tensor] = None) -> Tensor:
        with torch.no_grad():
            out = self.content_model.model(
                input_ids=content_ids,
                attention_mask=attention_mask,
            )
        return out.last_hidden_state

    def _context_forward(self, context_ids: Tensor,
                         attention_mask: Optional[Tensor] = None) -> Tensor:
        with torch.no_grad():
            out = self.context_model.model(
                input_ids=context_ids,
                attention_mask=attention_mask,
            )
        return out.last_hidden_state

    def forward(self, content_ids: Tensor, context_ids: Tensor,
                content_attention_mask: Optional[Tensor] = None,
                context_attention_mask: Optional[Tensor] = None,
                labels: Optional[Tensor] = None,
                use_context: bool = True,
                ) -> dict[str, Tensor]:
        """
        content_ids:  [B, T_c]   tokenized with content tokenizer
        context_ids:  [B, T_ctx] tokenized with context tokenizer (!)
        labels:       [B, T_c]   next-token targets (content tokenizer IDs)
        use_context:  if False, skip cross-attn gate — pure content model
        """

        # 1. Content: frozen content model
        content_x = self._content_forward(content_ids, content_attention_mask)

        if use_context:
            # 2. Context: frozen context model
            context_x = self._context_forward(context_ids, context_attention_mask)

            # Context padding mask (for cross-attention)
            ctx_pad = None
            if context_attention_mask is not None:
                ctx_pad = ~context_attention_mask.bool()

            # 3. Cross-attention gate
            gated = self.cross_gate(content_x, context_x, ctx_pad)
        else:
            gated = content_x

        # 4. LM head (from content model)
        logits = self.content_model.lm_head(gated)

        result: dict[str, Tensor] = {"logits": logits}
        if labels is not None:
            shift_logits = logits[:, :-1].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = nn.CrossEntropyLoss(ignore_index=-100)(
                shift_logits.view(-1, self.vocab_size), shift_labels.view(-1))
            result["loss"] = loss
        return result

    def trainable_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    @torch.no_grad()
    def gate_ablation_test(self, content_ids, context_ids,
                           content_attention_mask=None,
                           context_attention_mask=None,
                           labels=None) -> dict:
        """Measure how much the context stream contributes.

        Runs two forward passes:
          1. Gate active → loss_with_context
          2. Gate disabled (all zeros) → loss_without_context

        Returns: {loss_with, loss_without, contribution_pct}
        """
        self.eval()

        # Pass 1: gate active
        out = self.forward(content_ids, context_ids,
                          content_attention_mask, context_attention_mask, labels)
        loss_with = out["loss"].item()

        # Pass 2: gate disabled → use_context=False (no cross-attn, pure content)
        out = self.forward(content_ids, context_ids,
                          content_attention_mask, context_attention_mask,
                          labels, use_context=False)
        loss_without = out["loss"].item()

        contribution = (loss_without - loss_with) / max(loss_without, 1e-8) * 100

        return {
            "loss_with_context": loss_with,
            "loss_without_context": loss_without,
            "contribution_pct": contribution,
        }

    def _content_forward_cached(
        self, content_ids: Tensor,
        attention_mask: Optional[Tensor] = None,
        past_key_values: Optional[tuple] = None,
        use_cache: bool = True,
    ):
        """Content model forward with optional KV cache. Returns (hidden, past_key_values)."""
        with torch.no_grad():
            kwargs = {"input_ids": content_ids}
            if attention_mask is not None:
                kwargs["attention_mask"] = attention_mask
            if past_key_values is not None:
                kwargs["past_key_values"] = past_key_values
            if use_cache:
                kwargs["use_cache"] = True
            out = self.content_model.model(**kwargs)
        return out.last_hidden_state, getattr(out, "past_key_values", None)

    @torch.no_grad()
    def generate_kv(self, context_ids: Tensor, content_ids: Tensor,
                    max_new_tokens: int = 256, temperature: float = 0.7,
                    content_mask: Optional[Tensor] = None,
                    context_mask: Optional[Tensor] = None,
                    top_k: int = 0, top_p: float = 0.0,
                    record_gates: bool = False) -> tuple:
        """Generate with KV-cache: context encoded once, content incremental.

        Returns (token_ids, [gate_values]) if record_gates else (token_ids, []).
        """
        device = content_ids.device
        pad_id = self.content_tokenizer.pad_token_id or self.content_tokenizer.eos_token_id

        def sample(logits: Tensor) -> int:
            logits = logits / max(temperature, 1e-8)
            if top_k > 0:
                vals, idx = torch.topk(logits, min(top_k, logits.size(-1)))
                mask = torch.full_like(logits, float("-inf"))
                mask.scatter_(0, idx, vals)
                logits = torch.where(torch.isfinite(mask), mask, logits)
            if top_p > 0:
                sorted_logits, sorted_idx = torch.sort(logits, descending=True)
                cum_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
                cutoff = (cum_probs > top_p).nonzero(as_tuple=True)
                if len(cutoff[0]) > 0:
                    sorted_logits[cutoff[0][0] + 1:] = float("-inf")
                    logits = sorted_logits.scatter(0, sorted_idx, sorted_logits).clamp(min=float("-inf"))
            probs = torch.softmax(logits, dim=-1)
            return torch.multinomial(probs, 1).item()

        # ── 1. Encode context once, cache cross-attn K/V ────────────────
        context_x = self._context_forward(context_ids, context_mask)
        ctx_pad = None
        if context_mask is not None:
            ctx_pad = ~context_mask.bool()
        self.cross_gate.cache_context(context_x, ctx_pad)

        # ── 2. Encode full content once, get past_key_values ───────────
        content_hidden, past_kv = self._content_forward_cached(
            content_ids,
            attention_mask=None if content_mask is None else ~content_mask,
            use_cache=True,
        )

        gate_means: list[float] = []

        # ── 3. Gate + sample first token ───────────────────────────────
        last_hidden = content_hidden[:, -1:]  # [B, 1, D]
        gated = self.cross_gate.forward_step(last_hidden, record_gate=record_gates)
        if record_gates:
            gate_means.append(self.cross_gate._last_gate.mean().item())
        logits = self.content_model.lm_head(gated)[0, -1]
        tok = sample(logits)

        generated = [tok]
        next_ids = torch.tensor([[tok]], device=device, dtype=torch.long)

        # ── 4. Incremental: one token at a time ─────────────────────────
        for _ in range(max_new_tokens - 1):
            content_hidden, past_kv = self._content_forward_cached(
                next_ids,
                past_key_values=past_kv,
                use_cache=True,
            )
            last_hidden = content_hidden[:, -1:]  # [B, 1, D]
            gated = self.cross_gate.forward_step(last_hidden, record_gate=record_gates)
            if record_gates:
                gate_means.append(self.cross_gate._last_gate.mean().item())
            logits = self.content_model.lm_head(gated)[0, -1]
            tok = sample(logits)

            if tok == pad_id:
                break
            generated.append(tok)
            next_ids = torch.tensor([[tok]], device=device, dtype=torch.long)

        self.cross_gate.clear_cache()
        if record_gates:
            return generated, gate_means
        return generated, []

    def generate(self, context_text: str, content_text: str,
                 max_new_tokens: int = 512, temperature: float = 0.7) -> str:
        """Generate from context + content text strings."""

        ctx_enc = self.context_tokenizer(
            context_text, max_length=self.max_context_len,
            truncation=True, return_tensors="pt")
        cnt_enc = self.content_tokenizer(
            content_text, max_length=4096 - max_new_tokens,
            truncation=True, return_tensors="pt")

        ctx_ids = ctx_enc["input_ids"].to(self.cross_gate.q_proj.weight.device)
        cnt_ids = cnt_enc["input_ids"].to(self.cross_gate.q_proj.weight.device)

        pad_id = self.content_tokenizer.pad_token_id or self.content_tokenizer.eos_token_id

        generated: list[int] = []
        for _ in range(max_new_tokens):
            PAD = 4096
            if cnt_ids.size(1) >= PAD:
                padded = cnt_ids[:, -PAD:]
                seq_len = PAD
            else:
                pads = torch.full((1, PAD - cnt_ids.size(1)), pad_id,
                                  dtype=torch.long, device=cnt_ids.device)
                padded = torch.cat([cnt_ids, pads], dim=1)
                seq_len = cnt_ids.size(1)
            cont_mask = padded != pad_id

            out = self(
                content_ids=padded,
                context_ids=ctx_ids,
                content_attention_mask=cont_mask,
            )
            last = min(seq_len, PAD) - 1
            logits = out["logits"][0, last] / max(temperature, 1e-8)
            logits = torch.nan_to_num(logits, nan=-1e9, posinf=1e9, neginf=-1e9)
            probs = torch.softmax(logits, dim=-1).clamp(min=1e-12)
            tok = torch.multinomial(probs, 1).item()
            if tok == self.content_tokenizer.eos_token_id or tok == pad_id:
                break
            generated.append(tok)
            cnt_ids = torch.cat([cnt_ids, torch.tensor([[tok]], device=cnt_ids.device)], dim=1)

        return self.content_tokenizer.decode(generated, skip_special_tokens=True)