def handle(common, src, tgt, srcm, rsm, minfreq=False, vsize=False):
vcbc, nwordf = ldvocab(common,
minf=minfreq,
omit_vsize=vsize,
vanilla=False)
if src == common:
src_indices = None
else:
vcbw, nword = ldvocab(src,
minf=minfreq,
omit_vsize=vsize,
vanilla=False)
vcbw = reverse_dict(vcbw)
src_indices = torch.tensor(
[vcbc.get(vcbw[i], 0) for i in range(nword)], dtype=torch.long)
if tgt == common:
tgt_indices = None
else:
vcbw, nword = ldvocab(tgt,
minf=minfreq,
omit_vsize=vsize,
vanilla=False)
vcbw = reverse_dict(vcbw)
tgt_indices = torch.tensor(
[vcbc.get(vcbw[i], 0) for i in range(nword)], dtype=torch.long)
mymodel = NMT(cnfg.isize, nwordf, nwordf, cnfg.nlayer, cnfg.ff_hsize,
cnfg.drop, cnfg.attn_drop, cnfg.share_emb, cnfg.nhead,
cache_len_default, cnfg.attn_hsize, cnfg.norm_output,
cnfg.bindDecoderEmb, cnfg.forbidden_indexes)
mymodel = load_model_cpu(srcm, mymodel)
mymodel.update_vocab(src_indices=src_indices, tgt_indices=tgt_indices)
save_model(mymodel, rsm, sub_module=False, h5args=h5zipargs)
def __init__(self, modelfs, fvocab_i, fvocab_t, cnfg, minbsize=1, expand_for_mulgpu=True, bsize=64, maxpad=16, maxpart=4, maxtoken=1536, minfreq = False, vsize = False): vcbi, nwordi = ldvocab(fvocab_i, minfreq, vsize) vcbt, nwordt = ldvocab(fvocab_t, minfreq, vsize) self.vcbi, self.vcbt = vcbi, reverse_dict(vcbt) if expand_for_mulgpu: self.bsize = bsize * minbsize self.maxtoken = maxtoken * minbsize else: self.bsize = bsize self.maxtoken = maxtoken self.maxpad = maxpad self.maxpart = maxpart self.minbsize = minbsize if isinstance(modelfs, (list, tuple)): models = [] for modelf in modelfs: tmp = NMT(cnfg.isize, nwordi, nwordt, cnfg.nlayer, cnfg.ff_hsize, cnfg.drop, cnfg.attn_drop, cnfg.share_emb, cnfg.nhead, cache_len_default, cnfg.attn_hsize, cnfg.norm_output, cnfg.bindDecoderEmb, cnfg.forbidden_indexes) tmp = load_model_cpu(modelf, tmp) tmp.apply(load_fixing) models.append(tmp) model = Ensemble(models) else: model = NMT(cnfg.isize, nwordi, nwordt, cnfg.nlayer, cnfg.ff_hsize, cnfg.drop, cnfg.attn_drop, cnfg.share_emb, cnfg.nhead, cache_len_default, cnfg.attn_hsize, cnfg.norm_output, cnfg.bindDecoderEmb, cnfg.forbidden_indexes) model = load_model_cpu(modelfs, model) model.apply(load_fixing) model.eval() self.use_cuda, self.cuda_device, cuda_devices, self.multi_gpu = parse_cuda_decode(cnfg.use_cuda, cnfg.gpuid, cnfg.multi_gpu_decoding) if self.use_cuda: model.to(self.cuda_device) if self.multi_gpu: model = DataParallelMT(model, device_ids=cuda_devices, output_device=self.cuda_device.index, host_replicate=True, gather_output=False) self.use_amp = cnfg.use_amp and self.use_cuda self.beam_size = cnfg.beam_size self.length_penalty = cnfg.length_penalty self.net = model
def handle(srcf, rsf, typ, rlist):
td = h5py.File(cnfg.dev_data, "r")
nwordi = int(td["nwordi"][:][0])
nwordt = int(td["nwordt"][:][0])
td.close()
_tmpm = NMT(cnfg.isize, nwordi, nwordt, cnfg.nlayer, cnfg.ff_hsize,
cnfg.drop, cnfg.attn_drop, cnfg.share_emb, cnfg.nhead,
cnfg.cache_len, cnfg.attn_hsize, cnfg.norm_output,
cnfg.bindDecoderEmb, cnfg.forbidden_indexes)
_tmpm = load_model_cpu(srcf, _tmpm)
if typ == "enc":
_tmpm.enc.nets = ModuleList(remove_layers(list(_tmpm.enc.nets), rlist))
elif typ == "dec":
_tmpm.dec.nets = ModuleList(remove_layers(list(_tmpm.dec.nets), rlist))
save_model(_tmpm, rsf, False)
def load_fixing(module):
if hasattr(module, "fix_load"):
module.fix_load()
td = h5py.File(cnfg.dev_data, "r")
ntest = td["ndata"][:].item()
nword = td["nword"][:].tolist()
nwordi, nwordt = nword[0], nword[-1]
if len(sys.argv) == 2:
mymodel = NMT(cnfg.isize, nwordi, nwordt, cnfg.nlayer, cnfg.ff_hsize,
cnfg.drop, cnfg.attn_drop, cnfg.share_emb, cnfg.nhead,
cache_len_default, cnfg.attn_hsize, cnfg.norm_output,
cnfg.bindDecoderEmb, cnfg.forbidden_indexes)
mymodel = load_model_cpu(sys.argv[1], mymodel)
mymodel.apply(load_fixing)
else:
models = []
for modelf in sys.argv[1:]:
tmp = NMT(cnfg.isize, nwordi, nwordt, cnfg.nlayer, cnfg.ff_hsize,
cnfg.drop, cnfg.attn_drop, cnfg.share_emb, cnfg.nhead,
cache_len_default, cnfg.attn_hsize, cnfg.norm_output,
cnfg.bindDecoderEmb, cnfg.forbidden_indexes)
tmp = load_model_cpu(modelf, tmp)
tmp.apply(load_fixing)
return rs, cwd
def reverse_dict(din):
rs = {}
for k, v in din.items():
rs[v] = k
return rs
td = h5py.File(cnfg.test_data, "r")
nwordi = int(td["nwordi"][:][0])
td.close()
vcbt, nwordt = ldvocab(sys.argv[2])
vcbt = reverse_dict(vcbt)
mymodel = NMT(cnfg.isize, nwordi, nwordt, cnfg.nlayer, cnfg.ff_hsize,
cnfg.drop, cnfg.attn_drop, cnfg.share_emb, cnfg.nhead,
cnfg.cache_len, cnfg.attn_hsize)
mymodel.load_state_dict(torch.load(sys.argv[1], map_location='cpu'))
initmodel = NMT(cnfg.isize, nwordi, nwordt, cnfg.nlayer, cnfg.ff_hsize,
cnfg.drop, cnfg.attn_drop, cnfg.share_emb, cnfg.nhead,
cnfg.cache_len, cnfg.attn_hsize)
print(initmodel.enc.pemb.w.data.equal(initmodel.dec.pemb.w.data))
print(initmodel.enc.pemb.w.data.equal(mymodel.enc.pemb.w.data))
print(initmodel.enc.pemb.w.data.equal(mymodel.dec.pemb.w.data))
