""" 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)