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

605ddb0 verified 7 months ago

4.11 kB

"""
GLADIUS v2.0 — Tool Cortex
Tools as embeddings in the same vector space as vocabulary.
Tool activation via cosine similarity threshold — no JSON, no parsing.
SLA2 hybrid attention over tool registry at every layer (via kernel).
STUB VERSION — tool registry exists but no real tools are connected.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

from .config import KernelConfig


class ToolCortex(nn.Module):
    """
    Tool understanding via shared embedding space.
    Tools live in the same manifold as tokens. When a hidden state
    is close enough to a tool embedding (cosine sim > threshold),
    the tool activates. No special syntax needed.
    argmax_tool S(tool | hidden_state) where S = cosine_similarity
    """

    def __init__(self, config: KernelConfig):
        super().__init__()
        self.config = config

        # Tool embeddings: same dimension as token embeddings
        self.tool_embeddings = nn.Parameter(
            torch.randn(config.max_tools, config.hidden_dim) * 0.02
        )

        # Tool activation gate (refines raw cosine similarity)
        self.activation_gate = nn.Sequential(
            nn.Linear(config.hidden_dim * 2, config.hidden_dim // 2),
            nn.SiLU(),
            nn.Linear(config.hidden_dim // 2, 1),
            nn.Sigmoid(),
        )

        # Tool result projection (maps tool output back to hidden space)
        self.result_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)

        # Registry metadata (not learned — set at runtime)
        self.register_buffer('tool_active', torch.zeros(config.max_tools, dtype=torch.bool))
        self.num_registered = 0

    def register_grid_tools(self):
        """Register ARC grid manipulation tools (rotate, flip, fill, etc.)."""
        grid_tool_names = ['rotate', 'flip_h', 'flip_v', 'fill', 'crop', 'tile', 'recolor', 'overlay']
        for i, name in enumerate(grid_tool_names):
            if i < self.config.max_tools:
                with torch.no_grad():
                    self.tool_embeddings.data[i] = torch.randn(self.config.hidden_dim) * 0.01
                    self.tool_active[i] = True
                    self.num_registered += 1

    def register_tool(self, tool_id: int, description_embedding: torch.Tensor):
        """
        Register a tool by initializing its embedding.
        In the full system, description_embedding comes from encoding
        the tool's natural language description through the shared embeddings.
        """
        with torch.no_grad():
            self.tool_embeddings.data[tool_id] = description_embedding
            self.tool_active[tool_id] = True
            self.num_registered += 1

    def check_activation(self, hidden: torch.Tensor) -> torch.Tensor | None:
        """
        Check if any tool should activate based on hidden state similarity.
        Args:
            hidden: (batch, seq_len, hidden_dim)
        Returns:
            tool_contribution: (batch, seq_len, hidden_dim) or None
        """
        if self.num_registered == 0:
            return None

        # Pool hidden state
        pooled = hidden.mean(dim=1)  # (B, D)

        # Cosine similarity with all tool embeddings
        pooled_norm = F.normalize(pooled, dim=-1)
        tools_norm = F.normalize(self.tool_embeddings, dim=-1)
        similarities = torch.matmul(pooled_norm, tools_norm.T)  # (B, max_tools)

        # Mask inactive tools
        similarities = similarities.masked_fill(~self.tool_active.unsqueeze(0), -1.0)

        # Find best matching tool
        best_sim, best_tool = similarities.max(dim=-1)  # (B,)

        # Check threshold
        if best_sim.max().item() < self.config.tool_activation_threshold:
            return None

        # Tool activated — in full system, this would invoke the tool
        # STUB: return the tool embedding as contribution
        tool_embed = self.tool_embeddings[best_tool]  # (B, D)
        contribution = self.result_proj(tool_embed)

        # Broadcast across sequence
        return contribution.unsqueeze(1).expand_as(hidden)

	"""
	GLADIUS v2.0 — Tool Cortex

	Tools as embeddings in the same vector space as vocabulary.
	Tool activation via cosine similarity threshold — no JSON, no parsing.
	SLA2 hybrid attention over tool registry at every layer (via kernel).

	STUB VERSION — tool registry exists but no real tools are connected.
	"""

	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	from .config import KernelConfig


	class ToolCortex(nn.Module):
	"""
	Tool understanding via shared embedding space.

	Tools live in the same manifold as tokens. When a hidden state
	is close enough to a tool embedding (cosine sim > threshold),
	the tool activates. No special syntax needed.

	argmax_tool S(tool \| hidden_state) where S = cosine_similarity
	"""

	def __init__(self, config: KernelConfig):
	super().__init__()
	self.config = config

	# Tool embeddings: same dimension as token embeddings
	self.tool_embeddings = nn.Parameter(
	torch.randn(config.max_tools, config.hidden_dim) * 0.02
	)

	# Tool activation gate (refines raw cosine similarity)
	self.activation_gate = nn.Sequential(
	nn.Linear(config.hidden_dim * 2, config.hidden_dim // 2),
	nn.SiLU(),
	nn.Linear(config.hidden_dim // 2, 1),
	nn.Sigmoid(),
	)

	# Tool result projection (maps tool output back to hidden space)
	self.result_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)

	# Registry metadata (not learned — set at runtime)
	self.register_buffer('tool_active', torch.zeros(config.max_tools, dtype=torch.bool))
	self.num_registered = 0

	def register_grid_tools(self):
	"""Register ARC grid manipulation tools (rotate, flip, fill, etc.)."""
	grid_tool_names = ['rotate', 'flip_h', 'flip_v', 'fill', 'crop', 'tile', 'recolor', 'overlay']
	for i, name in enumerate(grid_tool_names):
	if i < self.config.max_tools:
	with torch.no_grad():
	self.tool_embeddings.data[i] = torch.randn(self.config.hidden_dim) * 0.01
	self.tool_active[i] = True
	self.num_registered += 1

	def register_tool(self, tool_id: int, description_embedding: torch.Tensor):
	"""
	Register a tool by initializing its embedding.

	In the full system, description_embedding comes from encoding
	the tool's natural language description through the shared embeddings.
	"""
	with torch.no_grad():
	self.tool_embeddings.data[tool_id] = description_embedding
	self.tool_active[tool_id] = True
	self.num_registered += 1

	def check_activation(self, hidden: torch.Tensor) -> torch.Tensor \| None:
	"""
	Check if any tool should activate based on hidden state similarity.

	Args:
	hidden: (batch, seq_len, hidden_dim)
	Returns:
	tool_contribution: (batch, seq_len, hidden_dim) or None
	"""
	if self.num_registered == 0:
	return None

	# Pool hidden state
	pooled = hidden.mean(dim=1) # (B, D)

	# Cosine similarity with all tool embeddings
	pooled_norm = F.normalize(pooled, dim=-1)
	tools_norm = F.normalize(self.tool_embeddings, dim=-1)
	similarities = torch.matmul(pooled_norm, tools_norm.T) # (B, max_tools)

	# Mask inactive tools
	similarities = similarities.masked_fill(~self.tool_active.unsqueeze(0), -1.0)

	# Find best matching tool
	best_sim, best_tool = similarities.max(dim=-1) # (B,)

	# Check threshold
	if best_sim.max().item() < self.config.tool_activation_threshold:
	return None

	# Tool activated — in full system, this would invoke the tool
	# STUB: return the tool embedding as contribution
	tool_embed = self.tool_embeddings[best_tool] # (B, D)
	contribution = self.result_proj(tool_embed)

	# Broadcast across sequence
	return contribution.unsqueeze(1).expand_as(hidden)