xSense/utils.py at master · DefinitionModeling/xSense · GitHub

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
import torch
import torch.nn as nn
from torch import optim
import numpy as np
from model import DecoderRNN, SPINEModel, MaskGenerator
from constants import VOC_DEC_NUM, VOC_W2V_NUM, EOS_IDX
import torch.nn.functional as F

USE_CUDA = torch.cuda.is_available()
device = torch.device("cuda" if USE_CUDA else "cpu")
torch.manual_seed(200)

def build_model_optimizer(args):
    print('Building models ...')
    if args.run == "pretrain":
        spine = SPINEModel(pretrain=True)
        spine = spine.to(device)
        spine_optimizer = optim.Adam(filter(lambda p: p.requires_grad, spine.parameters()), lr=args.lr)

        return spine_optimizer, spine

    else:
        # load pretrained model
        spine = SPINEModel(pretrain=False) #TODO: embed is removed !
        sp_ckpt = torch.load('save/model/pretrained_spine.tar')
        spine.load_state_dict(sp_ckpt['spine'])

        # build model
        n_voc_dec = VOC_DEC_NUM + 4
        dec_emb = nn.Embedding(n_voc_dec, 300)
        decoder = DecoderRNN(dec_emb, args.hidden_size, n_voc_dec, args.n_layers, args.dropout)
        mask_generator = MaskGenerator(z_dim=1000, enc_dim=300, K=args.K)

        # use cuda
        dec_emb = dec_emb.to(device)
        decoder = decoder.to(device)
        spine = spine.to(device)
        mask_generator = mask_generator.to(device)

        # optimizer
        print('Building optimizers ...')
        decoder_optimizer = optim.Adam(filter(lambda p: p.requires_grad, \
                                        list(decoder.parameters())+list(mask_generator.parameters())), lr=args.lr)

        return decoder_optimizer, decoder, mask_generator, spine


def load_model(args):
    # build model
    spine = SPINEModel(pretrain=False)
    n_voc_dec = VOC_DEC_NUM + 4
    decoder = DecoderRNN(nn.Embedding(n_voc_dec, 300), args.hidden_size, n_voc_dec, args.n_layers)
    mask_generator = MaskGenerator(z_dim=1000, enc_dim=300, K=args.K)
    # load from ckpt
    checkpoint = torch.load(args.model_path)
    decoder.load_state_dict(checkpoint['decoder'])
    spine.load_state_dict(checkpoint['spine'])
    mask_generator.load_state_dict(checkpoint['mask_gen'])
    decoder.eval()
    spine.eval()
    mask_generator.eval()

    # use cuda
    decoder = decoder.to(device)
    spine = spine.to(device)
    mask_generator = mask_generator.to(device)

    return decoder, spine, mask_generator


def get_mask(text_len):
    max_len = torch.max(text_len).item()
    idxes = torch.arange(0, max_len, out=torch.LongTensor(max_len)).unsqueeze(0)
    idxes = idxes.to(device)
    text_mask = (idxes < text_len.unsqueeze(1)).detach() # (batch, text_len)

    return text_mask, max_len


def compute_sparsity(X):
    non_zeros = 1. * np.count_nonzero(X)
    total = X.size
    sparsity = 100. * (1 - (non_zeros)/total)
    return sparsity


def maskNLLLoss(inp, target, mask):
    nTotal = mask.sum().item()
    loss_all = F.nll_loss(inp, target, reduction='none')
    loss = loss_all.masked_select(mask).mean()

    return loss, nTotal


def id2sent(sent_ids, voc_dec):
    words = []
    for idx in sent_ids:
        idx = idx.item()
        if idx == EOS_IDX:
            break
        else:
            words.append(voc_dec.index2word[idx])

    return ' '.join(words)


"""
def clip_parameters(model, clip):
    if clip > 0:
        for x in model.parameters():
            x.data.clamp_(-clip, clip)
"""