"""Small output dataclasses compatible with HF-style access patterns."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Tuple


@dataclass
class StrataBertModelOutput:
    last_hidden_state: object
    pooler_output: Optional[object] = None
    hidden_states: Optional[Tuple[object, ...]] = None
    attentions: Optional[Tuple[object, ...]] = None
    ssm_states: Optional[Tuple[object, ...]] = None


@dataclass
class StrataBertSequenceClassifierOutput:
    loss: Optional[object] = None
    logits: Optional[object] = None
    hidden_states: Optional[Tuple[object, ...]] = None
    attentions: Optional[Tuple[object, ...]] = None


@dataclass
class StrataBertTokenClassifierOutput:
    loss: Optional[object] = None
    logits: Optional[object] = None
    hidden_states: Optional[Tuple[object, ...]] = None
    attentions: Optional[Tuple[object, ...]] = None


@dataclass
class StrataBertMaskedLMOutput:
    loss: Optional[object] = None
    logits: Optional[object] = None
    hidden_states: Optional[Tuple[object, ...]] = None
    attentions: Optional[Tuple[object, ...]] = None


@dataclass
class StrataBertPreTrainingOutput:
    loss: Optional[object] = None
    prediction_logits: Optional[object] = None
    rtd_logits: Optional[object] = None
    mlm_loss: Optional[object] = None
    rtd_loss: Optional[object] = None
    hidden_states: Optional[Tuple[object, ...]] = None
    attentions: Optional[Tuple[object, ...]] = None