"""Shared model output types.

Models return logits deliberately.  Applying sigmoid in the loss makes mixed
precision training less stable and makes accidental double-sigmoid bugs easy.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
from torch import Tensor


@dataclass
class TurnDetectionOutput:
    """Raw outputs produced by every turn detector.

    Auxiliary filler logits are optional because the endpoint task can be
    trained on corpora which do not carry filler annotations.
    """

    endpoint_logits: Tensor
    midfiller_logits: Tensor | None = None
    endfiller_logits: Tensor | None = None
    embedding: Tensor | None = None

    def probabilities(self) -> dict[str, Tensor]:
        """Return calibrated-at-the-caller sigmoid probabilities."""

        result = {"endpoint": torch.sigmoid(self.endpoint_logits)}
        if self.midfiller_logits is not None:
            result["midfiller"] = torch.sigmoid(self.midfiller_logits)
        if self.endfiller_logits is not None:
            result["endfiller"] = torch.sigmoid(self.endfiller_logits)
        return result