Spaces:
Running
Running
| import torch | |
| from torch.nn import functional as F, Parameter | |
| from torch.autograd import Variable | |
| from torch.nn.init import xavier_normal_, xavier_uniform_ | |
| from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence | |
| class Distmult(torch.nn.Module): | |
| def __init__(self, args, num_entities, num_relations): | |
| super(Distmult, self).__init__() | |
| if args.max_norm: | |
| self.emb_e = torch.nn.Embedding(num_entities, args.embedding_dim, max_norm=1.0) | |
| self.emb_rel = torch.nn.Embedding(num_relations, args.embedding_dim) | |
| else: | |
| self.emb_e = torch.nn.Embedding(num_entities, args.embedding_dim, padding_idx=None) | |
| self.emb_rel = torch.nn.Embedding(num_relations, args.embedding_dim, padding_idx=None) | |
| self.inp_drop = torch.nn.Dropout(args.input_drop) | |
| self.loss = torch.nn.CrossEntropyLoss() | |
| self.init() | |
| def init(self): | |
| xavier_normal_(self.emb_e.weight) | |
| xavier_normal_(self.emb_rel.weight) | |
| def score_sr(self, sub, rel, sigmoid = False): | |
| sub_emb = self.emb_e(sub).squeeze(dim=1) | |
| rel_emb = self.emb_rel(rel).squeeze(dim=1) | |
| #sub_emb = self.inp_drop(sub_emb) | |
| #rel_emb = self.inp_drop(rel_emb) | |
| pred = torch.mm(sub_emb*rel_emb, self.emb_e.weight.transpose(1,0)) | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_or(self, obj, rel, sigmoid = False): | |
| obj_emb = self.emb_e(obj).squeeze(dim=1) | |
| rel_emb = self.emb_rel(rel).squeeze(dim=1) | |
| #obj_emb = self.inp_drop(obj_emb) | |
| #rel_emb = self.inp_drop(rel_emb) | |
| pred = torch.mm(obj_emb*rel_emb, self.emb_e.weight.transpose(1,0)) | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def forward(self, sub_emb, rel_emb, mode='rhs', sigmoid=False): | |
| ''' | |
| When mode is 'rhs' we expect (s,r); for 'lhs', we expect (o,r) | |
| For distmult, computations for both modes are equivalent, so we do not need if-else block | |
| ''' | |
| sub_emb = self.inp_drop(sub_emb) | |
| rel_emb = self.inp_drop(rel_emb) | |
| pred = torch.mm(sub_emb*rel_emb, self.emb_e.weight.transpose(1,0)) | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_triples(self, sub, rel, obj, sigmoid=False): | |
| ''' | |
| Inputs - subject, relation, object | |
| Return - score | |
| ''' | |
| sub_emb = self.emb_e(sub).squeeze(dim=1) | |
| rel_emb = self.emb_rel(rel).squeeze(dim=1) | |
| obj_emb = self.emb_e(obj).squeeze(dim=1) | |
| pred = torch.sum(sub_emb*rel_emb*obj_emb, dim=-1) | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_emb(self, emb_s, emb_r, emb_o, sigmoid=False): | |
| ''' | |
| Inputs - embeddings of subject, relation, object | |
| Return - score | |
| ''' | |
| pred = torch.sum(emb_s*emb_r*emb_o, dim=-1) | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_triples_vec(self, sub, rel, obj, sigmoid=False): | |
| ''' | |
| Inputs - subject, relation, object | |
| Return - a vector score for the triple instead of reducing over the embedding dimension | |
| ''' | |
| sub_emb = self.emb_e(sub).squeeze(dim=1) | |
| rel_emb = self.emb_rel(rel).squeeze(dim=1) | |
| obj_emb = self.emb_e(obj).squeeze(dim=1) | |
| pred = sub_emb*rel_emb*obj_emb | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| class Complex(torch.nn.Module): | |
| def __init__(self, args, num_entities, num_relations): | |
| super(Complex, self).__init__() | |
| if args.max_norm: | |
| self.emb_e = torch.nn.Embedding(num_entities, 2*args.embedding_dim, max_norm=1.0) | |
| self.emb_rel = torch.nn.Embedding(num_relations, 2*args.embedding_dim) | |
| else: | |
| self.emb_e = torch.nn.Embedding(num_entities, 2*args.embedding_dim, padding_idx=None) | |
| self.emb_rel = torch.nn.Embedding(num_relations, 2*args.embedding_dim, padding_idx=None) | |
| self.inp_drop = torch.nn.Dropout(args.input_drop) | |
| self.loss = torch.nn.CrossEntropyLoss() | |
| self.init() | |
| def init(self): | |
| xavier_normal_(self.emb_e.weight) | |
| xavier_normal_(self.emb_rel.weight) | |
| def score_sr(self, sub, rel, sigmoid = False): | |
| sub_emb = self.emb_e(sub).squeeze(dim=1) | |
| rel_emb = self.emb_rel(rel).squeeze(dim=1) | |
| s_real, s_img = torch.chunk(rel_emb, 2, dim=-1) | |
| rel_real, rel_img = torch.chunk(sub_emb, 2, dim=-1) | |
| emb_e_real, emb_e_img = torch.chunk(self.emb_e.weight, 2, dim=-1) | |
| realo_realreal = s_real*rel_real | |
| realo_imgimg = s_img*rel_img | |
| realo = realo_realreal - realo_imgimg | |
| real = torch.mm(realo, emb_e_real.transpose(1,0)) | |
| imgo_realimg = s_real*rel_img | |
| imgo_imgreal = s_img*rel_real | |
| imgo = imgo_realimg + imgo_imgreal | |
| img = torch.mm(imgo, emb_e_img.transpose(1,0)) | |
| pred = real + img | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_or(self, obj, rel, sigmoid = False): | |
| obj_emb = self.emb_e(obj).squeeze(dim=1) | |
| rel_emb = self.emb_rel(rel).squeeze(dim=1) | |
| rel_real, rel_img = torch.chunk(rel_emb, 2, dim=-1) | |
| o_real, o_img = torch.chunk(obj_emb, 2, dim=-1) | |
| emb_e_real, emb_e_img = torch.chunk(self.emb_e.weight, 2, dim=-1) | |
| #rel_real = self.inp_drop(rel_real) | |
| #rel_img = self.inp_drop(rel_img) | |
| #o_real = self.inp_drop(o_real) | |
| #o_img = self.inp_drop(o_img) | |
| # complex space bilinear product (equivalent to HolE) | |
| # realrealreal = torch.mm(rel_real*o_real, emb_e_real.transpose(1,0)) | |
| # realimgimg = torch.mm(rel_img*o_img, emb_e_real.transpose(1,0)) | |
| # imgrealimg = torch.mm(rel_real*o_img, emb_e_img.transpose(1,0)) | |
| # imgimgreal = torch.mm(rel_img*o_real, emb_e_img.transpose(1,0)) | |
| # pred = realrealreal + realimgimg + imgrealimg - imgimgreal | |
| reals_realreal = rel_real*o_real | |
| reals_imgimg = rel_img*o_img | |
| reals = reals_realreal + reals_imgimg | |
| real = torch.mm(reals, emb_e_real.transpose(1,0)) | |
| imgs_realimg = rel_real*o_img | |
| imgs_imgreal = rel_img*o_real | |
| imgs = imgs_realimg - imgs_imgreal | |
| img = torch.mm(imgs, emb_e_img.transpose(1,0)) | |
| pred = real + img | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def forward(self, sub_emb, rel_emb, mode='rhs', sigmoid=False): | |
| ''' | |
| When mode is 'rhs' we expect (s,r); for 'lhs', we expect (o,r) | |
| ''' | |
| if mode == 'lhs': | |
| rel_real, rel_img = torch.chunk(rel_emb, 2, dim=-1) | |
| o_real, o_img = torch.chunk(sub_emb, 2, dim=-1) | |
| emb_e_real, emb_e_img = torch.chunk(self.emb_e.weight, 2, dim=-1) | |
| rel_real = self.inp_drop(rel_real) | |
| rel_img = self.inp_drop(rel_img) | |
| o_real = self.inp_drop(o_real) | |
| o_img = self.inp_drop(o_img) | |
| reals_realreal = rel_real*o_real | |
| reals_imgimg = rel_img*o_img | |
| reals = reals_realreal + reals_imgimg | |
| real = torch.mm(reals, emb_e_real.transpose(1,0)) | |
| imgs_realimg = rel_real*o_img | |
| imgs_imgreal = rel_img*o_real | |
| imgs = imgs_realimg - imgs_imgreal | |
| img = torch.mm(imgs, emb_e_img.transpose(1,0)) | |
| pred = real + img | |
| else: | |
| s_real, s_img = torch.chunk(rel_emb, 2, dim=-1) | |
| rel_real, rel_img = torch.chunk(sub_emb, 2, dim=-1) | |
| emb_e_real, emb_e_img = torch.chunk(self.emb_e.weight, 2, dim=-1) | |
| s_real = self.inp_drop(s_real) | |
| s_img = self.inp_drop(s_img) | |
| rel_real = self.inp_drop(rel_real) | |
| rel_img = self.inp_drop(rel_img) | |
| realo_realreal = s_real*rel_real | |
| realo_imgimg = s_img*rel_img | |
| realo = realo_realreal - realo_imgimg | |
| real = torch.mm(realo, emb_e_real.transpose(1,0)) | |
| imgo_realimg = s_real*rel_img | |
| imgo_imgreal = s_img*rel_real | |
| imgo = imgo_realimg + imgo_imgreal | |
| img = torch.mm(imgo, emb_e_img.transpose(1,0)) | |
| pred = real + img | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_triples(self, sub, rel, obj, sigmoid=False): | |
| ''' | |
| Inputs - subject, relation, object | |
| Return - score | |
| ''' | |
| sub_emb = self.emb_e(sub).squeeze(dim=1) | |
| rel_emb = self.emb_rel(rel).squeeze(dim=1) | |
| obj_emb = self.emb_e(obj).squeeze(dim=1) | |
| s_real, s_img = torch.chunk(sub_emb, 2, dim=-1) | |
| rel_real, rel_img = torch.chunk(rel_emb, 2, dim=-1) | |
| o_real, o_img = torch.chunk(obj_emb, 2, dim=-1) | |
| realrealreal = torch.sum(s_real*rel_real*o_real, dim=-1) | |
| realimgimg = torch.sum(s_real*rel_img*o_img, axis=-1) | |
| imgrealimg = torch.sum(s_img*rel_real*o_img, axis=-1) | |
| imgimgreal = torch.sum(s_img*rel_img*o_real, axis=-1) | |
| pred = realrealreal + realimgimg + imgrealimg - imgimgreal | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_emb(self, emb_s, emb_r, emb_o, sigmoid=False): | |
| ''' | |
| Inputs - embeddings of subject, relation, object | |
| Return - score | |
| ''' | |
| s_real, s_img = torch.chunk(emb_s, 2, dim=-1) | |
| rel_real, rel_img = torch.chunk(emb_r, 2, dim=-1) | |
| o_real, o_img = torch.chunk(emb_o, 2, dim=-1) | |
| realrealreal = torch.sum(s_real*rel_real*o_real, dim=-1) | |
| realimgimg = torch.sum(s_real*rel_img*o_img, axis=-1) | |
| imgrealimg = torch.sum(s_img*rel_real*o_img, axis=-1) | |
| imgimgreal = torch.sum(s_img*rel_img*o_real, axis=-1) | |
| pred = realrealreal + realimgimg + imgrealimg - imgimgreal | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_triples_vec(self, sub, rel, obj, sigmoid=False): | |
| ''' | |
| Inputs - subject, relation, object | |
| Return - a vector score for the triple instead of reducing over the embedding dimension | |
| ''' | |
| sub_emb = self.emb_e(sub).squeeze(dim=1) | |
| rel_emb = self.emb_rel(rel).squeeze(dim=1) | |
| obj_emb = self.emb_e(obj).squeeze(dim=1) | |
| s_real, s_img = torch.chunk(sub_emb, 2, dim=-1) | |
| rel_real, rel_img = torch.chunk(rel_emb, 2, dim=-1) | |
| o_real, o_img = torch.chunk(obj_emb, 2, dim=-1) | |
| realrealreal = s_real*rel_real*o_real | |
| realimgimg = s_real*rel_img*o_img | |
| imgrealimg = s_img*rel_real*o_img | |
| imgimgreal = s_img*rel_img*o_real | |
| pred = realrealreal + realimgimg + imgrealimg - imgimgreal | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| class Conve(torch.nn.Module): | |
| #Too slow !!!! | |
| def __init__(self, args, num_entities, num_relations): | |
| super(Conve, self).__init__() | |
| if args.max_norm: | |
| self.emb_e = torch.nn.Embedding(num_entities, args.embedding_dim, max_norm=1.0) | |
| self.emb_rel = torch.nn.Embedding(num_relations, args.embedding_dim) | |
| else: | |
| self.emb_e = torch.nn.Embedding(num_entities, args.embedding_dim, padding_idx=None) | |
| self.emb_rel = torch.nn.Embedding(num_relations, args.embedding_dim, padding_idx=None) | |
| self.inp_drop = torch.nn.Dropout(args.input_drop) | |
| self.hidden_drop = torch.nn.Dropout(args.hidden_drop) | |
| self.feature_drop = torch.nn.Dropout2d(args.feat_drop) | |
| self.embedding_dim = args.embedding_dim #default is 200 | |
| self.num_filters = args.num_filters # default is 32 | |
| self.kernel_size = args.kernel_size # default is 3 | |
| self.stack_width = args.stack_width # default is 20 | |
| self.stack_height = args.embedding_dim // self.stack_width | |
| self.bn0 = torch.nn.BatchNorm2d(1) | |
| self.bn1 = torch.nn.BatchNorm2d(self.num_filters) | |
| self.bn2 = torch.nn.BatchNorm1d(args.embedding_dim) | |
| self.conv1 = torch.nn.Conv2d(1, out_channels=self.num_filters, | |
| kernel_size=(self.kernel_size, self.kernel_size), | |
| stride=1, padding=0, bias=args.use_bias) | |
| #self.conv1 = torch.nn.Conv2d(1, 32, (3, 3), 1, 0, bias=args.use_bias) # <-- default | |
| flat_sz_h = int(2*self.stack_width) - self.kernel_size + 1 | |
| flat_sz_w = self.stack_height - self.kernel_size + 1 | |
| self.flat_sz = flat_sz_h*flat_sz_w*self.num_filters | |
| self.fc = torch.nn.Linear(self.flat_sz, args.embedding_dim) | |
| self.register_parameter('b', Parameter(torch.zeros(num_entities))) | |
| self.loss = torch.nn.CrossEntropyLoss() | |
| self.init() | |
| def init(self): | |
| xavier_normal_(self.emb_e.weight) | |
| xavier_normal_(self.emb_rel.weight) | |
| def concat(self, e1_embed, rel_embed, form='plain'): | |
| if form == 'plain': | |
| e1_embed = e1_embed. view(-1, 1, self.stack_width, self.stack_height) | |
| rel_embed = rel_embed.view(-1, 1, self.stack_width, self.stack_height) | |
| stack_inp = torch.cat([e1_embed, rel_embed], 2) | |
| elif form == 'alternate': | |
| e1_embed = e1_embed. view(-1, 1, self.embedding_dim) | |
| rel_embed = rel_embed.view(-1, 1, self.embedding_dim) | |
| stack_inp = torch.cat([e1_embed, rel_embed], 1) | |
| stack_inp = torch.transpose(stack_inp, 2, 1).reshape((-1, 1, 2*self.stack_width, self.stack_height)) | |
| else: raise NotImplementedError | |
| return stack_inp | |
| def conve_architecture(self, sub_emb, rel_emb): | |
| stacked_inputs = self.concat(sub_emb, rel_emb) | |
| stacked_inputs = self.bn0(stacked_inputs) | |
| x = self.inp_drop(stacked_inputs) | |
| x = self.conv1(x) | |
| x = self.bn1(x) | |
| x = F.relu(x) | |
| x = self.feature_drop(x) | |
| #x = x.view(x.shape[0], -1) | |
| x = x.view(-1, self.flat_sz) | |
| x = self.fc(x) | |
| x = self.hidden_drop(x) | |
| x = self.bn2(x) | |
| x = F.relu(x) | |
| return x | |
| def score_sr(self, sub, rel, sigmoid = False): | |
| sub_emb = self.emb_e(sub) | |
| rel_emb = self.emb_rel(rel) | |
| x = self.conve_architecture(sub_emb, rel_emb) | |
| pred = torch.mm(x, self.emb_e.weight.transpose(1,0)) | |
| pred += self.b.expand_as(pred) | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_or(self, obj, rel, sigmoid = False): | |
| obj_emb = self.emb_e(obj) | |
| rel_emb = self.emb_rel(rel) | |
| x = self.conve_architecture(obj_emb, rel_emb) | |
| pred = torch.mm(x, self.emb_e.weight.transpose(1,0)) | |
| pred += self.b.expand_as(pred) | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def forward(self, sub_emb, rel_emb, mode='rhs', sigmoid=False): | |
| ''' | |
| When mode is 'rhs' we expect (s,r); for 'lhs', we expect (o,r) | |
| For conve, computations for both modes are equivalent, so we do not need if-else block | |
| ''' | |
| x = self.conve_architecture(sub_emb, rel_emb) | |
| pred = torch.mm(x, self.emb_e.weight.transpose(1,0)) | |
| pred += self.b.expand_as(pred) | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_triples(self, sub, rel, obj, sigmoid=False): | |
| ''' | |
| Inputs - subject, relation, object | |
| Return - score | |
| ''' | |
| sub_emb = self.emb_e(sub) | |
| rel_emb = self.emb_rel(rel) | |
| obj_emb = self.emb_e(obj) | |
| x = self.conve_architecture(sub_emb, rel_emb) | |
| pred = torch.mm(x, obj_emb.transpose(1,0)) | |
| #print(pred.shape) | |
| pred += self.b[obj].expand_as(pred) #taking the bias value for object embedding | |
| # above works fine for single input triples; | |
| # but if input is batch of triples, then this is a matrix of (num_trip x num_trip) where diagonal is scores | |
| # so use torch.diagonal() after calling this function | |
| pred = torch.diagonal(pred) | |
| # or could have used : pred= torch.sum(x*obj_emb, dim=-1) | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_emb(self, emb_s, emb_r, emb_o, sigmoid=False): | |
| ''' | |
| Inputs - embeddings of subject, relation, object | |
| Return - score | |
| ''' | |
| x = self.conve_architecture(emb_s, emb_r) | |
| pred = torch.mm(x, emb_o.transpose(1,0)) | |
| #pred += self.b[obj].expand_as(pred) #taking the bias value for object embedding - don't know which obj | |
| # above works fine for single input triples; | |
| # but if input is batch of triples, then this is a matrix of (num_trip x num_trip) where diagonal is scores | |
| # so use torch.diagonal() after calling this function | |
| pred = torch.diagonal(pred) | |
| # or could have used : pred= torch.sum(x*obj_emb, dim=-1) | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred | |
| def score_triples_vec(self, sub, rel, obj, sigmoid=False): | |
| ''' | |
| Inputs - subject, relation, object | |
| Return - a vector score for the triple instead of reducing over the embedding dimension | |
| ''' | |
| sub_emb = self.emb_e(sub) | |
| rel_emb = self.emb_rel(rel) | |
| obj_emb = self.emb_e(obj) | |
| x = self.conve_architecture(sub_emb, rel_emb) | |
| #pred = torch.mm(x, obj_emb.transpose(1,0)) | |
| pred = x*obj_emb | |
| #print(pred.shape, self.b[obj].shape) #shapes are [7,200] and [7] | |
| #pred += self.b[obj].expand_as(pred) #taking the bias value for object embedding - can't add scalar to vector | |
| #pred = sub_emb*rel_emb*obj_emb | |
| if sigmoid: | |
| pred = torch.sigmoid(pred) | |
| return pred |