tinymind-native-colab-handoff/bundle/model/self_assessment_core.py

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn

from .layers import RMSNorm


@dataclass
class SelfAssessmentCoreConfig:
    dim: int
    inner_dim: int | None = None
    recursion_steps: int = 2
    residual_scale: float = 0.15
    dropout: float = 0.0


class SelfAssessmentCore(nn.Module):
    """Recursive hidden-state assessor for confidence, conflict, and evidence.

    The module is deliberately compact: it does not decide final truth by
    itself. It estimates uncertainty signals from the current hidden state and
    feeds a bounded correction back into the residual stream.
    """

    def __init__(self, cfg: SelfAssessmentCoreConfig):
        super().__init__()
        self.cfg = cfg
        inner = int(cfg.inner_dim or cfg.dim * 2)
        self.norm = RMSNorm(cfg.dim)
        self.evidence_proj = nn.Sequential(nn.Linear(cfg.dim, inner), nn.SiLU(), nn.Linear(inner, 1))
        self.conflict_proj = nn.Sequential(nn.Linear(cfg.dim, inner), nn.SiLU(), nn.Linear(inner, 1))
        self.correction = nn.Sequential(nn.Linear(cfg.dim + 3, inner), nn.SiLU(), nn.Dropout(cfg.dropout), nn.Linear(inner, cfg.dim))
        self.logic_gate = nn.Linear(cfg.dim, cfg.dim)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        h = self.norm(x)
        evidence = torch.sigmoid(self.evidence_proj(h))
        conflict = torch.sigmoid(self.conflict_proj(h))
        uncertainty = torch.clamp(conflict * (1.0 - evidence), 0.0, 1.0)
        confidence = torch.clamp(evidence * (1.0 - conflict), 0.0, 1.0)

        logic = h
        steps = max(1, int(self.cfg.recursion_steps))
        for _ in range(steps):
            # Bounded nested logic: each pass can refine, never explode.
            logic = torch.tanh(self.logic_gate(logic))

        features = torch.cat([logic, confidence, uncertainty, conflict], dim=-1)
        delta = torch.tanh(self.correction(features))
        scale = float(max(0.0, min(self.cfg.residual_scale, 1.0)))
        y = x + scale * (1.0 - uncertainty) * delta
        report = {
            "confidence": confidence,
            "uncertainty": uncertainty,
            "conflict": conflict,
            "evidence": evidence,
        }
        return y, report