WYRM kernel v27 FINAL — live training version (443.7M base + Synthase + PUP + GaussianHead)

605ddb0 verified 7 months ago

11.8 kB

"""
GLADIUS v2.0 — The Kernel
This is the core. Everything flows through here.
Input → Embed → Memory Read → Time Stamp → Transformer Layers →
  Router → Specialists → Tool Check → Modulate → Decode →
  Memory Write → Cognition Check → Output
Every decision is argmax S(x | context).
"""

import torch
import torch.nn as nn
import time as time_module

from .config import KernelConfig
from .embeddings import SharedEmbeddings
from .attention import TransformerLayer, RMSNorm
from .memory import ThreeTemperatureMemory
from .temporal import TimeEngine
from .cognition import CognitionLoop
from .modulator import Modulator
from .tools import ToolCortex
from .router import NexusRouter
from .senses import SensoryCortex, VisionConfig, AudioConfig


class GladiusKernel(nn.Module):
    """
    The GLADIUS Kernel.
    Not a model. Not a wrapper. A kernel.
    Memory manages persistence. Cognition schedules thinking.
    Time provides awareness. Modulator controls voice.
    Tool Cortex provides hands. Specialists run ON this kernel.
    """

    def __init__(self, config: KernelConfig,
                 vision_config: VisionConfig | None = None,
                 audio_config: AudioConfig | None = None):
        super().__init__()
        self.config = config

        # === Core Components ===
        self.embeddings = SharedEmbeddings(config)
        self.memory = ThreeTemperatureMemory(config)
        self.time_engine = TimeEngine(config)
        self.cognition = CognitionLoop(config)
        self.modulator = Modulator(config)
        self.tool_cortex = ToolCortex(config)
        self.router = NexusRouter(config)

        # === Sensory Cortex (optional — additive, never required) ===
        self.has_senses = vision_config is not None or audio_config is not None
        if self.has_senses:
            self.senses = SensoryCortex(config, vision_config, audio_config)
        else:
            self.senses = None

        # === Transformer Backbone ===
        self.layers = nn.ModuleList([
            TransformerLayer(config, layer_idx=i)
            for i in range(config.num_layers)
        ])

        # === Final Norm ===
        self.final_norm = RMSNorm(config.hidden_dim)

        # === Causal Mask Cache ===
        self.register_buffer(
            'causal_mask',
            torch.tril(torch.ones(config.max_seq_len, config.max_seq_len))
                .unsqueeze(0).unsqueeze(0)  # (1, 1, S, S)
        )

        # Print parameter count on init
        self._report_params()

    def _report_params(self):
        total = sum(p.numel() for p in self.parameters())
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        print(f"GLADIUS Kernel initialized: {total:,} params ({trainable:,} trainable)")
        print(f"  Memory: {total * 4 / 1024 / 1024:.1f} MB (float32)")

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        timestamp: float | torch.Tensor | None = None,
        images: torch.Tensor | None = None,
        audio: torch.Tensor | None = None,
    ) -> dict:
        """
        Full forward pass through the kernel.
        Args:
            input_ids: (batch, seq_len) token IDs — can be None for pure sensory input
            timestamp: Unix timestamp (or None for current time)
            images: (batch, C, H, W) pixel values [0, 1] — vision input
            audio: (batch, 1, n_mels, n_frames) mel spectrogram — audio input
        Returns:
            dict with:
                logits: (batch, seq_len, vocab_size) — modulated output logits
                silence: (batch, 1) — silence gate value
                mode: CognitiveMode — current cognitive mode
                importance: (batch, seq_len, 1) — memory importance scores
                modality_mask: (batch, seq_len) — 0=text, 1=vision, 2=audio (if multimodal)
        """
        # 1. Embed text tokens (if provided)
        text_embeds = None
        if input_ids is not None:
            B, S = input_ids.shape
            text_embeds = self.embeddings.embed(input_ids)  # (B, S, D)
        
        # 2. Sensory integration
        modality_mask = None
        if self.has_senses and (images is not None or audio is not None):
            x, modality_mask = self.senses(
                text_embeds=text_embeds,
                images=images,
                audio=audio,
            )
            B = x.shape[0]
            S = x.shape[1]
        elif text_embeds is not None:
            x = text_embeds
            B, S = x.shape[0], x.shape[1]
        else:
            raise ValueError("Must provide input_ids, images, or audio")

        # 2. Memory read (hot memory context + warm adapter)
        x = self.memory.read(x)

        # 3. Temporal encoding (additive input + stored for output gating)
        time_embed = None
        if timestamp is not None:
            if isinstance(timestamp, (int, float)):
                timestamp = torch.tensor([timestamp] * B, dtype=torch.float32)
            time_embed = self.time_engine(timestamp)  # (B, D)
            x = x + time_embed.unsqueeze(1)  # Broadcast across seq_len

        # 4. Transformer layers with causal mask
        # Dynamic mask — sequence may exceed max_seq_len when multimodal
        if S <= self.config.max_seq_len:
            mask = self.causal_mask[:, :, :S, :S]
        else:
            mask = torch.tril(torch.ones(1, 1, S, S, device=x.device))
        for layer in self.layers:
            x = layer(x, mask=mask)

        # 5. Final norm
        x = self.final_norm(x)

        # 6. Tool check
        tool_result = self.tool_cortex.check_activation(x)
        if tool_result is not None:
            x = x + tool_result

        # 7. Modulate and produce logits (time gates the output)
        logits, silence, pixel_output = self.modulator(x, self.embeddings.output_head, temporal_embedding=time_embed)

        # 8. Memory write
        importance = self.memory.write(x)

        # 9. Cognition heartbeat — returns (mode, cognitive_state, mode_probs)
        mode, cognitive_state, mode_probs = self.cognition.heartbeat(x)

        # 10. Consolidation check
        if self.cognition.should_consolidate():
            self.memory.consolidate()

        # Record event in time engine
        self.time_engine.record_event()

        return {
            'logits': logits,
            'silence': silence,
            'pixel_output': pixel_output, # Added pixel_output
            'mode': mode,
            'importance': importance,
            'modality_mask': modality_mask,
            'cognitive_state': cognitive_state,
            'mode_probs': mode_probs,
        }

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_tokens: int = 100,
        temperature: float = 1.0,
        top_k: int = 50,
        timestamp: float | None = None,
    ) -> torch.Tensor:
        """
        Autoregressive generation.
        Args:
            input_ids: (1, seq_len) — prompt tokens
            max_tokens: maximum tokens to generate
            temperature: sampling temperature (1.0 = neutral)
            top_k: top-k sampling (0 = greedy)
            timestamp: Unix timestamp
        Returns:
            (1, seq_len + generated) — full sequence
        """
        self.eval()

        if timestamp is None:
            timestamp = time_module.time()

        for _ in range(max_tokens):
            # Truncate to max_seq_len
            context = input_ids[:, -self.config.max_seq_len:]

            # Forward pass
            result = self.forward(context, timestamp=timestamp)
            logits = result['logits'][:, -1, :]  # Last position
            silence = result['silence']

            # Silence gate check
            if silence.item() > self.config.silence_threshold:
                break

            # Temperature scaling
            if temperature != 1.0:
                logits = logits / temperature

            # Top-k sampling
            if top_k > 0:
                topk_logits, topk_indices = logits.topk(top_k, dim=-1)
                probs = torch.softmax(topk_logits, dim=-1)
                sampled_idx = torch.multinomial(probs, 1)
                next_token = topk_indices.gather(-1, sampled_idx)
            else:
                next_token = logits.argmax(dim=-1, keepdim=True)

            # Append
            input_ids = torch.cat([input_ids, next_token], dim=1)

            # Stop on EOS
            if next_token.item() == self.config.eos_token_id:
                break

        return input_ids

    def save_checkpoint(self, path: str):
        """Save full kernel state."""
        torch.save({
            'model_state_dict': self.state_dict(),
            'config': self.config,
        }, path)
        # Also save warm memory separately (for restart survival)
        self.memory.checkpoint(path + '.warm')

    @classmethod
    def load_checkpoint(cls, path: str, map_location: str | None = None) -> 'GladiusKernel':
        """Load kernel from checkpoint."""
        data = torch.load(path, map_location=map_location, weights_only=False)
        
        import dataclasses
        cfg_raw = data['config']
        if isinstance(cfg_raw, dict):
            config_dict = cfg_raw
        elif dataclasses.is_dataclass(cfg_raw) and not isinstance(cfg_raw, type):
            config_dict = dataclasses.asdict(cfg_raw)
        else:
            config_dict = dict(cfg_raw)
        
        # Ensure cold_embedding_dim matches hidden_dim
        if 'cold_embedding_dim' not in config_dict or config_dict['cold_embedding_dim'] != config_dict['hidden_dim']:
            config_dict['cold_embedding_dim'] = config_dict['hidden_dim']
        
        # Filter out keys not in KernelConfig (e.g., 'clock_mode' from older checkpoints)
        valid_fields = {f.name for f in dataclasses.fields(KernelConfig)}
        extra_keys = {k for k in config_dict if k not in valid_fields}
        filtered_config = {k: v for k, v in config_dict.items() if k in valid_fields}
        
        # Handle dtype serialization (torch.dtype may come back as string or int)
        if 'dtype' in filtered_config:
            dtype_val = filtered_config['dtype']
            if isinstance(dtype_val, str):
                filtered_config['dtype'] = getattr(torch, dtype_val.replace('torch.', ''), torch.float32)
            elif not isinstance(dtype_val, torch.dtype):
                filtered_config['dtype'] = torch.float32
        
        # Preserve clock_mode as attribute after construction if needed
        clock_mode = config_dict.get('clock_mode', 'continuous')
        
        config_from_checkpoint = KernelConfig(**filtered_config)
        
        # Set clock_mode as attribute (used by TimeEngine)
        if hasattr(config_from_checkpoint, 'clock_mode'):
            config_from_checkpoint.clock_mode = clock_mode
        elif clock_mode != 'continuous':
            config_from_checkpoint.clock_mode = clock_mode
        
        # Infer missing config from state dict shapes (backward compat)
        sd = data['model_state_dict']
        if 'tool_cortex.tool_embeddings' in sd:
            actual_max_tools = sd['tool_cortex.tool_embeddings'].shape[0]
            if config_from_checkpoint.max_tools != actual_max_tools:
                config_from_checkpoint.max_tools = actual_max_tools
        
        kernel = cls(config_from_checkpoint)
        # Load state_dict, ignoring missing keys (for newly added heads like pixel_head)
        kernel.load_state_dict(data['model_state_dict'], strict=False)
        # Restore warm memory if available
        try:
            kernel.memory.restore(path + '.warm')
        except FileNotFoundError:
            pass
        return kernel

	"""
	GLADIUS v2.0 — The Kernel

	This is the core. Everything flows through here.

	Input → Embed → Memory Read → Time Stamp → Transformer Layers →
	Router → Specialists → Tool Check → Modulate → Decode →
	Memory Write → Cognition Check → Output

	Every decision is argmax S(x \| context).
	"""

	import torch
	import torch.nn as nn
	import time as time_module

	from .config import KernelConfig
	from .embeddings import SharedEmbeddings
	from .attention import TransformerLayer, RMSNorm
	from .memory import ThreeTemperatureMemory
	from .temporal import TimeEngine
	from .cognition import CognitionLoop
	from .modulator import Modulator
	from .tools import ToolCortex
	from .router import NexusRouter
	from .senses import SensoryCortex, VisionConfig, AudioConfig


	class GladiusKernel(nn.Module):
	"""
	The GLADIUS Kernel.

	Not a model. Not a wrapper. A kernel.
	Memory manages persistence. Cognition schedules thinking.
	Time provides awareness. Modulator controls voice.
	Tool Cortex provides hands. Specialists run ON this kernel.
	"""

	def __init__(self, config: KernelConfig,
	vision_config: VisionConfig \| None = None,
	audio_config: AudioConfig \| None = None):
	super().__init__()
	self.config = config

	# === Core Components ===
	self.embeddings = SharedEmbeddings(config)
	self.memory = ThreeTemperatureMemory(config)
	self.time_engine = TimeEngine(config)
	self.cognition = CognitionLoop(config)
	self.modulator = Modulator(config)
	self.tool_cortex = ToolCortex(config)
	self.router = NexusRouter(config)

	# === Sensory Cortex (optional — additive, never required) ===
	self.has_senses = vision_config is not None or audio_config is not None
	if self.has_senses:
	self.senses = SensoryCortex(config, vision_config, audio_config)
	else:
	self.senses = None

	# === Transformer Backbone ===
	self.layers = nn.ModuleList([
	TransformerLayer(config, layer_idx=i)
	for i in range(config.num_layers)
	])

	# === Final Norm ===
	self.final_norm = RMSNorm(config.hidden_dim)

	# === Causal Mask Cache ===
	self.register_buffer(
	'causal_mask',
	torch.tril(torch.ones(config.max_seq_len, config.max_seq_len))
	.unsqueeze(0).unsqueeze(0) # (1, 1, S, S)
	)

	# Print parameter count on init
	self._report_params()

	def _report_params(self):
	total = sum(p.numel() for p in self.parameters())
	trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
	print(f"GLADIUS Kernel initialized: {total:,} params ({trainable:,} trainable)")
	print(f" Memory: {total * 4 / 1024 / 1024:.1f} MB (float32)")

	def forward(
	self,
	input_ids: torch.Tensor \| None = None,
	timestamp: float \| torch.Tensor \| None = None,
	images: torch.Tensor \| None = None,
	audio: torch.Tensor \| None = None,
	) -> dict:
	"""
	Full forward pass through the kernel.

	Args:
	input_ids: (batch, seq_len) token IDs — can be None for pure sensory input
	timestamp: Unix timestamp (or None for current time)
	images: (batch, C, H, W) pixel values [0, 1] — vision input
	audio: (batch, 1, n_mels, n_frames) mel spectrogram — audio input

	Returns:
	dict with:
	logits: (batch, seq_len, vocab_size) — modulated output logits
	silence: (batch, 1) — silence gate value
	mode: CognitiveMode — current cognitive mode
	importance: (batch, seq_len, 1) — memory importance scores
	modality_mask: (batch, seq_len) — 0=text, 1=vision, 2=audio (if multimodal)
	"""
	# 1. Embed text tokens (if provided)
	text_embeds = None
	if input_ids is not None:
	B, S = input_ids.shape
	text_embeds = self.embeddings.embed(input_ids) # (B, S, D)

	# 2. Sensory integration
	modality_mask = None
	if self.has_senses and (images is not None or audio is not None):
	x, modality_mask = self.senses(
	text_embeds=text_embeds,
	images=images,
	audio=audio,
	)
	B = x.shape[0]
	S = x.shape[1]
	elif text_embeds is not None:
	x = text_embeds
	B, S = x.shape[0], x.shape[1]
	else:
	raise ValueError("Must provide input_ids, images, or audio")

	# 2. Memory read (hot memory context + warm adapter)
	x = self.memory.read(x)

	# 3. Temporal encoding (additive input + stored for output gating)
	time_embed = None
	if timestamp is not None:
	if isinstance(timestamp, (int, float)):
	timestamp = torch.tensor([timestamp] * B, dtype=torch.float32)
	time_embed = self.time_engine(timestamp) # (B, D)
	x = x + time_embed.unsqueeze(1) # Broadcast across seq_len

	# 4. Transformer layers with causal mask
	# Dynamic mask — sequence may exceed max_seq_len when multimodal
	if S <= self.config.max_seq_len:
	mask = self.causal_mask[:, :, :S, :S]
	else:
	mask = torch.tril(torch.ones(1, 1, S, S, device=x.device))
	for layer in self.layers:
	x = layer(x, mask=mask)

	# 5. Final norm
	x = self.final_norm(x)

	# 6. Tool check
	tool_result = self.tool_cortex.check_activation(x)
	if tool_result is not None:
	x = x + tool_result

	# 7. Modulate and produce logits (time gates the output)
	logits, silence, pixel_output = self.modulator(x, self.embeddings.output_head, temporal_embedding=time_embed)

	# 8. Memory write
	importance = self.memory.write(x)

	# 9. Cognition heartbeat — returns (mode, cognitive_state, mode_probs)
	mode, cognitive_state, mode_probs = self.cognition.heartbeat(x)

	# 10. Consolidation check
	if self.cognition.should_consolidate():
	self.memory.consolidate()

	# Record event in time engine
	self.time_engine.record_event()

	return {
	'logits': logits,
	'silence': silence,
	'pixel_output': pixel_output, # Added pixel_output
	'mode': mode,
	'importance': importance,
	'modality_mask': modality_mask,
	'cognitive_state': cognitive_state,
	'mode_probs': mode_probs,
	}

	@torch.no_grad()
	def generate(
	self,
	input_ids: torch.Tensor,
	max_tokens: int = 100,
	temperature: float = 1.0,
	top_k: int = 50,
	timestamp: float \| None = None,
	) -> torch.Tensor:
	"""
	Autoregressive generation.

	Args:
	input_ids: (1, seq_len) — prompt tokens
	max_tokens: maximum tokens to generate
	temperature: sampling temperature (1.0 = neutral)
	top_k: top-k sampling (0 = greedy)
	timestamp: Unix timestamp

	Returns:
	(1, seq_len + generated) — full sequence
	"""
	self.eval()

	if timestamp is None:
	timestamp = time_module.time()

	for _ in range(max_tokens):
	# Truncate to max_seq_len
	context = input_ids[:, -self.config.max_seq_len:]

	# Forward pass
	result = self.forward(context, timestamp=timestamp)
	logits = result['logits'][:, -1, :] # Last position
	silence = result['silence']

	# Silence gate check
	if silence.item() > self.config.silence_threshold:
	break

	# Temperature scaling
	if temperature != 1.0:
	logits = logits / temperature

	# Top-k sampling
	if top_k > 0:
	topk_logits, topk_indices = logits.topk(top_k, dim=-1)
	probs = torch.softmax(topk_logits, dim=-1)
	sampled_idx = torch.multinomial(probs, 1)
	next_token = topk_indices.gather(-1, sampled_idx)
	else:
	next_token = logits.argmax(dim=-1, keepdim=True)

	# Append
	input_ids = torch.cat([input_ids, next_token], dim=1)

	# Stop on EOS
	if next_token.item() == self.config.eos_token_id:
	break

	return input_ids

	def save_checkpoint(self, path: str):
	"""Save full kernel state."""
	torch.save({
	'model_state_dict': self.state_dict(),
	'config': self.config,
	}, path)
	# Also save warm memory separately (for restart survival)
	self.memory.checkpoint(path + '.warm')

	@classmethod
	def load_checkpoint(cls, path: str, map_location: str \| None = None) -> 'GladiusKernel':
	"""Load kernel from checkpoint."""
	data = torch.load(path, map_location=map_location, weights_only=False)

	import dataclasses
	cfg_raw = data['config']
	if isinstance(cfg_raw, dict):
	config_dict = cfg_raw
	elif dataclasses.is_dataclass(cfg_raw) and not isinstance(cfg_raw, type):
	config_dict = dataclasses.asdict(cfg_raw)
	else:
	config_dict = dict(cfg_raw)

	# Ensure cold_embedding_dim matches hidden_dim
	if 'cold_embedding_dim' not in config_dict or config_dict['cold_embedding_dim'] != config_dict['hidden_dim']:
	config_dict['cold_embedding_dim'] = config_dict['hidden_dim']

	# Filter out keys not in KernelConfig (e.g., 'clock_mode' from older checkpoints)
	valid_fields = {f.name for f in dataclasses.fields(KernelConfig)}
	extra_keys = {k for k in config_dict if k not in valid_fields}
	filtered_config = {k: v for k, v in config_dict.items() if k in valid_fields}

	# Handle dtype serialization (torch.dtype may come back as string or int)
	if 'dtype' in filtered_config:
	dtype_val = filtered_config['dtype']
	if isinstance(dtype_val, str):
	filtered_config['dtype'] = getattr(torch, dtype_val.replace('torch.', ''), torch.float32)
	elif not isinstance(dtype_val, torch.dtype):
	filtered_config['dtype'] = torch.float32

	# Preserve clock_mode as attribute after construction if needed
	clock_mode = config_dict.get('clock_mode', 'continuous')

	config_from_checkpoint = KernelConfig(**filtered_config)

	# Set clock_mode as attribute (used by TimeEngine)
	if hasattr(config_from_checkpoint, 'clock_mode'):
	config_from_checkpoint.clock_mode = clock_mode
	elif clock_mode != 'continuous':
	config_from_checkpoint.clock_mode = clock_mode

	# Infer missing config from state dict shapes (backward compat)
	sd = data['model_state_dict']
	if 'tool_cortex.tool_embeddings' in sd:
	actual_max_tools = sd['tool_cortex.tool_embeddings'].shape[0]
	if config_from_checkpoint.max_tools != actual_max_tools:
	config_from_checkpoint.max_tools = actual_max_tools

	kernel = cls(config_from_checkpoint)
	# Load state_dict, ignoring missing keys (for newly added heads like pixel_head)
	kernel.load_state_dict(data['model_state_dict'], strict=False)
	# Restore warm memory if available
	try:
	kernel.memory.restore(path + '.warm')
	except FileNotFoundError:
	pass
	return kernel