File size: 2,402 Bytes
1502abe f70f582 1502abe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """Merlin-Agent modeling: Qwen3_5 text causal LM with per-layer quantum injection.
At each full-attention layer, a fixed quantum-derived direction (buffer q_k) is
added to that layer's output hidden state with an RMS-matched, alpha-scaled
magnitude. Injection lives in the model (buffers + hooks re-installed in
__init__), so it survives save/reload — not a bare, non-persisted hook.
alpha == 0 -> exact base model.
"""
import torch
try:
from transformers import Qwen3_5ForCausalLM
except ImportError: # not exported at top level in some transformers versions
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM
try:
from .configuration_merlin_agent import MerlinAgentConfig
except ImportError: # when loaded as flat remote code (trust_remote_code)
from configuration_merlin_agent import MerlinAgentConfig
def _inject(h: torch.Tensor, q: torch.Tensor, alpha: float) -> torch.Tensor:
if alpha == 0:
return h
q = q.to(dtype=h.dtype, device=h.device)
rms_h = h.pow(2).mean(dim=-1, keepdim=True).sqrt()
rms_q = q.pow(2).mean().sqrt().clamp_min(1e-6)
return h + alpha * (rms_h / rms_q) * q
class MerlinAgentForCausalLM(Qwen3_5ForCausalLM):
config_class = MerlinAgentConfig
def __init__(self, config):
super().__init__(config)
self._inj_layers = list(config.full_attention_layer_indices)
for k in range(len(self._inj_layers)):
self.register_buffer(f"q_{k}", torch.zeros(config.hidden_size), persistent=True)
self._install_injection_hooks()
def _install_injection_hooks(self):
for k, li in enumerate(self._inj_layers):
self.model.layers[li].register_forward_hook(self._make_hook(f"q_{k}"))
def _make_hook(self, qk: str):
def hook(module, args, output):
q = getattr(self, qk)
a = float(self.config.quantum_injection_alpha)
if isinstance(output, tuple):
return (_inject(output[0], q, a),) + tuple(output[1:])
return _inject(output, q, a)
return hook
def set_quantum_signatures(self, q_vectors):
"""Load the 8 projected quantum direction vectors (each shape [hidden_size])."""
assert len(q_vectors) == len(self._inj_layers)
for k, v in enumerate(q_vectors):
getattr(self, f"q_{k}").copy_(torch.as_tensor(v, dtype=torch.float32))
|