-
Notifications
You must be signed in to change notification settings - Fork 0
/
de_simple.py
117 lines (98 loc) · 6.13 KB
/
de_simple.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
# Copyright (c) 2018-present, Royal Bank of Canada.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from params import Params
from dataset import Dataset
class DE_SimplE(torch.nn.Module):
def __init__(self, dataset, params):
super(DE_SimplE, self).__init__()
self.dataset = dataset
self.params = params
self.ent_embs_h = nn.Embedding(dataset.numEnt(), params.s_emb_dim).cuda()
self.ent_embs_t = nn.Embedding(dataset.numEnt(), params.s_emb_dim).cuda()
self.rel_embs_f = nn.Embedding(dataset.numRel(), params.s_emb_dim+params.t_emb_dim).cuda()
self.rel_embs_i = nn.Embedding(dataset.numRel(), params.s_emb_dim+params.t_emb_dim).cuda()
self.create_time_embedds()
self.time_nl = torch.sin
nn.init.xavier_uniform_(self.ent_embs_h.weight)
nn.init.xavier_uniform_(self.ent_embs_t.weight)
nn.init.xavier_uniform_(self.rel_embs_f.weight)
nn.init.xavier_uniform_(self.rel_embs_i.weight)
def create_time_embedds(self):
# frequency embeddings for the entities
self.m_freq_h = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.m_freq_t = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.d_freq_h = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.d_freq_t = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.y_freq_h = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.y_freq_t = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
# phi embeddings for the entities
self.m_phi_h = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.m_phi_t = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.d_phi_h = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.d_phi_t = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.y_phi_h = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.y_phi_t = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
# frequency embeddings for the entities
self.m_amps_h = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.m_amps_t = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.d_amps_h = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.d_amps_t = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.y_amps_h = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
self.y_amps_t = nn.Embedding(self.dataset.numEnt(), self.params.t_emb_dim).cuda()
nn.init.xavier_uniform_(self.m_freq_h.weight)
nn.init.xavier_uniform_(self.d_freq_h.weight)
nn.init.xavier_uniform_(self.y_freq_h.weight)
nn.init.xavier_uniform_(self.m_freq_t.weight)
nn.init.xavier_uniform_(self.d_freq_t.weight)
nn.init.xavier_uniform_(self.y_freq_t.weight)
nn.init.xavier_uniform_(self.m_phi_h.weight)
nn.init.xavier_uniform_(self.d_phi_h.weight)
nn.init.xavier_uniform_(self.y_phi_h.weight)
nn.init.xavier_uniform_(self.m_phi_t.weight)
nn.init.xavier_uniform_(self.d_phi_t.weight)
nn.init.xavier_uniform_(self.y_phi_t.weight)
nn.init.xavier_uniform_(self.m_amps_h.weight)
nn.init.xavier_uniform_(self.d_amps_h.weight)
nn.init.xavier_uniform_(self.y_amps_h.weight)
nn.init.xavier_uniform_(self.m_amps_t.weight)
nn.init.xavier_uniform_(self.d_amps_t.weight)
nn.init.xavier_uniform_(self.y_amps_t.weight)
def get_time_embedd(self, entities, years, months, days, h_or_t):
if h_or_t == "head":
emb = self.y_amps_h(entities) * self.time_nl(self.y_freq_h(entities) * years + self.y_phi_h(entities))
emb += self.m_amps_h(entities) * self.time_nl(self.m_freq_h(entities) * months + self.m_phi_h(entities))
emb += self.d_amps_h(entities) * self.time_nl(self.d_freq_h(entities) * days + self.d_phi_h(entities))
else:
emb = self.y_amps_t(entities) * self.time_nl(self.y_freq_t(entities) * years + self.y_phi_t(entities))
emb += self.m_amps_t(entities) * self.time_nl(self.m_freq_t(entities) * months + self.m_phi_t(entities))
emb += self.d_amps_t(entities) * self.time_nl(self.d_freq_t(entities) * days + self.d_phi_t(entities))
return emb
def getEmbeddings(self, heads, rels, tails, years, months, days, intervals = None):
years = years.view(-1,1)
months = months.view(-1,1)
days = days.view(-1,1)
h_embs1 = self.ent_embs_h(heads)
r_embs1 = self.rel_embs_f(rels)
t_embs1 = self.ent_embs_t(tails)
h_embs2 = self.ent_embs_h(tails)
r_embs2 = self.rel_embs_i(rels)
t_embs2 = self.ent_embs_t(heads)
h_embs1 = torch.cat((h_embs1, self.get_time_embedd(heads, years, months, days, "head")), 1)
t_embs1 = torch.cat((t_embs1, self.get_time_embedd(tails, years, months, days, "tail")), 1)
h_embs2 = torch.cat((h_embs2, self.get_time_embedd(tails, years, months, days, "head")), 1)
t_embs2 = torch.cat((t_embs2, self.get_time_embedd(heads, years, months, days, "tail")), 1)
return h_embs1, r_embs1, t_embs1, h_embs2, r_embs2, t_embs2
def forward(self, heads, rels, tails, years, months, days):
h_embs1, r_embs1, t_embs1, h_embs2, r_embs2, t_embs2 = self.getEmbeddings(heads, rels, tails, years, months, days)
scores = ((h_embs1 * r_embs1) * t_embs1 + (h_embs2 * r_embs2) * t_embs2) / 2.0
scores = F.dropout(scores, p=self.params.dropout, training=self.training)
scores = torch.sum(scores, dim=1)
return scores