models/losses.py

from typing import Any, Tuple, Dict, Sequence, Optional

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

# This ID tells the code to "skip" or ignore certain data points (like padding)
IGNORE_LABEL_ID = -100


def s(x, epsilon=1e-30):
    """
    A helper function called 'StableMax' transformation.
    Instead of standard Softmax, this is more 'stable' for the GPU.
    If a number is negative, it squashes it; if positive, it slightly increases it.
    """
    return torch.where(
        x < 0,
        1 / (1 - x + epsilon),
        x + 1
    )


def log_stablemax(x, dim=-1):
    """
    Calculates the 'Log-Probability'.
    In AI, we prefer working with 'Logs' because adding small numbers is 
    much safer for computers than multiplying them (which can lead to 0).
    """
    s_x = s(x)
    return torch.log(s_x / torch.sum(s_x, dim=dim, keepdim=True))


def stablemax_cross_entropy(logits, labels, ignore_index: int = -100):
    """
    The main 'Comparison' function. 
    It compares the model's guess (logits) to the real answer (labels).
    The higher the score, the worse the model's guess was.
    """
    # 1. Convert the model's numbers into probabilities
    logprobs = log_stablemax(logits.to(torch.float64), dim=-1)

    # 2. figure out which cells actually have numbers (and aren't just empty padding)
    valid_mask = labels != ignore_index
    transformed_labels = torch.where(valid_mask, labels, 0)
    
    # 3. Pick out the probability for the correct answer
    prediction_logprobs = torch.gather(logprobs, index=transformed_labels.to(torch.long).unsqueeze(-1), dim=-1).squeeze(-1)

    # 4. Return the negative of that probability (our 'Punishment' score)
    return -torch.where(valid_mask, prediction_logprobs, 0)


def softmax_cross_entropy(logits, labels, ignore_index: int = -100):
    # Cast logits to f32
    # Flatten logits
    return F.cross_entropy(logits.to(torch.float32).view(-1, logits.shape[-1]), labels.to(torch.long).view(-1), ignore_index=ignore_index, reduction="none").view(labels.shape)


class ACTLossHead(nn.Module):
    """
    The 'Reasoning Manager'. 
    This class wraps the model and decides:
    1. Did the model solve it?
    2. Should it stop thinking?
    3. How many steps did it take?
    """
    def __init__(self, model: nn.Module, loss_type: str):
        super().__init__()
        self.model = model
        self.loss_fn = globals()[loss_type]
        
    def initial_carry(self, *args, **kwargs):
        """Prepares the model's 'Internal Memory' for a new puzzle."""
        return self.model.initial_carry(*args, **kwargs)  # type: ignore

    def forward(
        self,
        return_keys: Sequence[str],
        # Model args
        **model_kwargs,
    ) -> Tuple[Any, torch.Tensor, Dict[str, torch.Tensor], Optional[Dict[str, torch.Tensor]], torch.Tensor]:
        # Model logits
        # B x SeqLen x D
        new_carry, outputs = self.model(**model_kwargs)
        labels = new_carry.current_data["labels"]

        # 2. Check the results (without keeping track of 'learning math' here)
        with torch.no_grad():
            mask = labels != IGNORE_LABEL_ID
            loss_counts = mask.sum(-1)
            loss_divisor = loss_counts.clamp_min(1).unsqueeze(-1)

            # Did the model pick the right numbers?
            is_correct = mask & (torch.argmax(outputs["logits"], dim=-1) == labels)
            # Is the WHOLE puzzle correct?
            seq_is_correct = is_correct.sum(-1) == loss_counts
            
            # --- Store metrics for the human to read(When it halts) ---
            valid_metrics = new_carry.halted & (loss_counts > 0)
            metrics = {
                "count": valid_metrics.sum(),
                "accuracy":       torch.where(valid_metrics, (is_correct.to(torch.float32) / loss_divisor).sum(-1), 0).sum(),
                "exact_accuracy": (valid_metrics & seq_is_correct).sum(),
                "q_halt_accuracy": (valid_metrics & ((outputs["q_halt_logits"] >= 0) == seq_is_correct)).sum(),
                "steps":          torch.where(valid_metrics, new_carry.steps, 0).sum(),
            }

        # 3. Calculate the learning punishments (Losses)
        # lm_loss: Punishment for getting digits wrong
        lm_loss = (self.loss_fn(outputs["logits"], labels, ignore_index=IGNORE_LABEL_ID) / loss_divisor).sum()
        
        # q_halt_loss: Punishment for stopping too early (or too late!)
        q_halt_loss = F.binary_cross_entropy_with_logits(outputs["q_halt_logits"], seq_is_correct.to(outputs["q_halt_logits"].dtype), reduction="sum")

        metrics.update({
            "lm_loss": lm_loss.detach(),
            "q_halt_loss": q_halt_loss.detach(),
        })

        # 4. Q-Continue Loss: Bootstrapping (The 'Look-Ahead' logic)
        q_continue_loss = 0
        if "target_q_continue" in outputs:
            # Encourages the model to continue if thinking more will lead to a correct answer
            q_continue_loss = F.binary_cross_entropy_with_logits(outputs["q_continue_logits"], outputs["target_q_continue"], reduction="sum")
            metrics["q_continue_loss"] = q_continue_loss.detach()

        # Filter outputs for return
        detached_outputs = {k: outputs[k].detach() for k in return_keys if k in outputs}

        return new_carry, lm_loss + 0.5 * (q_halt_loss + q_continue_loss), metrics, detached_outputs, new_carry.halted.all()