import torch
from torch import nn
import torch.nn.functional as F


class TransE(nn.Module):
    def __init__(self, margin=1.0):
        super(TransE, self).__init__()
        self.margin = margin

    def forward(self, cui_0, cui_1, cui_2, re):
        pos = cui_0 + re - cui_1
        neg = cui_0 + re - cui_2
        return torch.mean(F.relu(self.margin + torch.norm(pos, p=2, dim=1) - torch.norm(neg, p=2, dim=1)))