ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """Production of outputs from current, persistent and abstract toroidal state.""" | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class ProductionHead(nn.Module): | |
| def __init__(self, d_model=256, vocab_size=32000): | |
| super().__init__() | |
| self.d_model, self.vocab_size = d_model, vocab_size | |
| self.decoder = nn.Sequential(nn.Linear(d_model, d_model), nn.GELU(), nn.Linear(d_model, vocab_size)) | |
| self.confidence_head = nn.Linear(d_model, 1) | |
| def produce(self, alpha, persistent_state=None, abstractions=None): | |
| state = alpha.mean(dim=0) if alpha.dim() > 1 else alpha | |
| if persistent_state is not None: | |
| state = state + (persistent_state.mean(dim=0) if persistent_state.dim() > 1 else persistent_state) | |
| if abstractions: | |
| reps = [] | |
| for a in abstractions: | |
| p = a["pattern"].to(state.device, state.dtype) | |
| if p.numel() == self.d_model: | |
| reps.append(p.reshape(self.d_model)) | |
| if reps: | |
| state = state + 0.1 * torch.stack(reps).mean(dim=0) | |
| logits = self.decoder(state).unsqueeze(0) | |
| confidence = torch.sigmoid(self.confidence_head(state)).reshape(1) | |
| return logits, confidence | |
| def sample(self, alpha, persistent_state=None, abstractions=None, temperature=1.0, top_k=50): | |
| if temperature <= 0: | |
| raise ValueError("temperature must be > 0") | |
| logits, confidence = self.produce(alpha, persistent_state, abstractions) | |
| logits = logits / temperature | |
| vocab = logits.shape[-1] | |
| if top_k is not None and top_k > 0 and top_k < vocab: | |
| values, indices = torch.topk(logits, top_k, dim=-1) | |
| logits = torch.full_like(logits, float("-inf")).scatter(-1, indices, values) | |
| probs = F.softmax(logits, dim=-1) | |
| token_id = torch.multinomial(probs.squeeze(0), 1).item() | |
| return token_id, logits, confidence | |