from enum import Enum
from typing import Annotated

import torch

EntityID = Annotated[int, "Entity ID"]
RelationID = Annotated[int, "Relation ID"]

PADDING_VALUE = -1  # used as a padding value in tensors of EntityID and RelationID

EntityIDTensor = Annotated[torch.Tensor, "Entity ID Tensor"]
RelationIDTensor = Annotated[torch.Tensor, "Relation ID Tensor"]
EmbeddingTensor = Annotated[torch.Tensor, "Embedding Tensor"]


class GrammaticalFunction(Enum):
    SUBJECT = "subject"
    RELATION = "relation"
    OBJECT = "object"
