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
vcbt, nwordt = ldvocab(sys.argv[3])
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,
cache_len_default, cnfg.attn_hsize, cnfg.norm_output,
cnfg.bindDecoderEmb, cnfg.forbidden_indexes, cnfg.num_layer_fwd)
mymodel = load_model_cpu(sys.argv[2], mymodel)
mymodel.apply(load_fixing)
mymodel.eval()
enc, trans, classifier = mymodel.enc, mymodel.dec.trans, mymodel.dec.classifier
use_cuda, cuda_device, cuda_devices, multi_gpu = parse_cuda_decode(
cnfg.use_cuda, cnfg.gpuid, cnfg.multi_gpu_decoding)
# do not support multi-gpu
multi_gpu, cuda_devices = False, None
use_amp = cnfg.use_amp and use_cuda
set_random_seed(cnfg.seed, use_cuda)
if cuda_device:
enc.to(cuda_device)
trans.to(cuda_device)
classifier.to(cuda_device)
ignore_ids = set(init_vocab.values())
src_grp = td["src"]
ens = "\n".encode("utf-8")
