File size: 2,545 Bytes
0523be0 | 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 | ###############################################################################
# 6. FULL COGNITHREAD MODEL + GGUF EXPORT
###############################################################################
class CogniThreadModel(nn.Module):
"""
Full CogniThread SSM architecture combining:
1. CogniThreadSSM (Mamba-style SSM with ZOH + parallel scan)
2. FuzzyHeuristicRetriever (Gaussian fuzzy membership)
3. LatentDiffusionSynthesizer (conditioned reverse diffusion)
4. AdaptiveGate (variance-constrained gating)
Forward:
h_t = SSM(x) # SSM trajectory
retrieved, w_f, mu = Fuzzy(h_t) # fuzzy retrieval
z_0 = Diffusion(h_t) # diffusion generation
y_final, var = Gate(retrieved, z_0, w_f) # adaptive gating
"""
def __init__(self, d_model: int, d_state: int = 16, n_heuristics: int = 64,
d_z: int = 32, T_diff: int = 10, d_conv: int = 4,
expand: int = 2, gamma: float = 1.0, top_k: int = 5, *,
rngs: nn.Rngs):
super().__init__()
self.d_model = d_model
self.d_state = d_state
self.d_z = d_z
self.ssm = CogniThreadSSM(d_model, d_state, d_conv, expand, rngs=rngs)
# State projection from SSM output dim to d_state for fuzzy/diffusion
self.state_proj = nn.Linear(d_model, d_state, rngs=rngs)
self.fuzzy = FuzzyHeuristicRetriever(d_state, n_heuristics, gamma, top_k, rngs=rngs)
self.diffusion = LatentDiffusionSynthesizer(d_z, d_state, T_diff, rngs=rngs)
# Project diffusion output to d_model
self.diff_proj = nn.Linear(d_z, d_model, rngs=rngs)
# Project fuzzy retrieved to d_model
self.fuzzy_proj = nn.Linear(d_state, d_model, rngs=rngs)
self.gate = AdaptiveGate(d_model, rngs=rngs)
def __call__(self, x, key=None):
if key is None:
key = jax.random.key(0)
# SSM forward
h_ssm = self.ssm(x) # (batch, seq_len, d_model)
# Project to state space for fuzzy/diffusion
h_t = self.state_proj(h_ssm) # (batch, seq_len, d_state)
# Fuzzy heuristic retrieval
retrieved, w_f, mu = self.fuzzy(h_t)
retrieved = self.fuzzy_proj(retrieved) # -> d_model
# Diffusion generation
z_0 = self.diffusion(h_t, key=key)
generated = self.diff_proj(z_0) # -> d_model
# Adaptive gating
y_final, var_y = self.gate(retrieved, generated, w_f)
return y_final, var_y, w_f, mu
print("[6/6] CogniThreadModel defined")
|