| |
| |
| |
| class CogniThreadV2(nn.Module): |
| """Full CogniThread + vocab head for next-token prediction.""" |
| def __init__(self, d_model, d_state=16, n_heuristics=64, d_z=32, |
| T_diff=10, d_conv=4, expand=2, vocab_size=256, *, |
| rngs: nn.Rngs): |
| super().__init__() |
| self.d_model = d_model |
| self.d_state = d_state |
| self.d_z = d_z |
| self.vocab_size = vocab_size |
| |
| self.embed = nn.Embed(vocab_size, d_model, rngs=rngs) |
| |
| self.ssm = CogniThreadSSM(d_model, d_state, d_conv, expand, rngs=rngs) |
| self.state_proj = nn.Linear(d_model, d_state, rngs=rngs) |
| self.fuzzy = FuzzyHeuristicRetriever(d_state, n_heuristics, 1.0, 5, rngs=rngs) |
| self.diffusion = LatentDiffusionSynthesizer(d_z, d_state, T_diff, rngs=rngs) |
| self.diff_proj = nn.Linear(d_z, d_model, rngs=rngs) |
| self.fuzzy_proj = nn.Linear(d_state, d_model, rngs=rngs) |
| self.gate = AdaptiveGate(d_model, rngs=rngs) |
| |
| self.lm_head = nn.Linear(d_model, vocab_size, rngs=rngs) |
|
|
| def __call__(self, input_ids, key=None): |
| """input_ids: (batch, seq_len) int tokens. Returns logits.""" |
| if key is None: key = jax.random.key(0) |
| x = self.embed(input_ids) |
| h_ssm = self.ssm(x) |
| h_t = self.state_proj(h_ssm) |
| retrieved, w_f, mu = self.fuzzy(h_t) |
| retrieved = self.fuzzy_proj(retrieved) |
| z_0 = self.diffusion(h_t, key=key) |
| generated = self.diff_proj(z_0) |
| y_final, var_y = self.gate(retrieved, generated, w_f) |
| logits = self.lm_head(y_final) |
| return logits |
|
|
| def generate(self, input_ids, max_new_tokens=64, key=None, temperature=0.8): |
| """Autoregressive text generation.""" |
| if key is None: key = jax.random.key(0) |
| batch, seq_len = input_ids.shape |
| generated = list(input_ids.T) |
| cur_ids = input_ids |
| for i in range(max_new_tokens): |
| key, sk = jax.random.split(key) |
| logits = self(cur_ids, key=sk) |
| next_logits = logits[:, -1, :] / max(temperature, 0.01) |
| next_token = jax.random.categorical(sk, next_logits, axis=-1) |
| generated.append(next_token) |
| cur_ids = jnp.concatenate([cur_ids[:, 1:], next_token[:, None]], axis=1) |
| return jnp.stack(generated, axis=1) |
|
|
