File size: 1,132 Bytes
35d483e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """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
|