from dataclasses import dataclass
from typing import Optional

import torch
from transformers.file_utils import ModelOutput


@dataclass
class RelationAwareClassificationModelOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    loss_cls: Optional[torch.FloatTensor] = None
    loss_span: Optional[torch.FloatTensor] = None
    class_logits: torch.FloatTensor = None
    span_logits: torch.FloatTensor = None