KG-Evaluation__Bachelor-Thesis__/tools/utils.py

import random
from enum import Enum

import numpy as np
import torch


def set_seed(seed: int = 42) -> None:
    """
    Set the random seed for reproducibility.

    Args:
        seed (int): The seed value to set for random number generation.
    """
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.default_rng(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


class Target(str, Enum):  # ↔ PyKEEN's LABEL_* constants
    HEAD = "head"
    TAIL = "tail"
    RELATION = "relation"