def __init__(self, params):
reloaded = torch.load(params.model_path, map_location='cpu')
# print(reloaded['dico_word2id']['while'])
# print(reloaded['dico_word2id']['return'])
# print(reloaded['dico_word2id']['if'])
# print(reloaded['encoder'].keys())
# print(reloaded['decoder'].keys())
reloaded['encoder'] = {(k[len('module.'):] if k.startswith('module.') else k): v for k, v in
reloaded['encoder'].items()}
assert 'decoder' in reloaded or (
'decoder_0' in reloaded and 'decoder_1' in reloaded)
if 'decoder' in reloaded:
decoders_names = ['decoder']
else:
decoders_names = ['decoder_0', 'decoder_1']
for decoder_name in decoders_names:
reloaded[decoder_name] = {(k[len('module.'):] if k.startswith('module.') else k): v for k, v in
reloaded[decoder_name].items()}
self.reloaded_params = AttrDict(reloaded['params'])
# build dictionary / update parameters
self.dico = Dictionary(
reloaded['dico_id2word'], reloaded['dico_word2id'], reloaded['dico_counts'])
assert self.reloaded_params.n_words == len(self.dico)
assert self.reloaded_params.bos_index == self.dico.index(BOS_WORD)
assert self.reloaded_params.eos_index == self.dico.index(EOS_WORD)
assert self.reloaded_params.pad_index == self.dico.index(PAD_WORD)
assert self.reloaded_params.unk_index == self.dico.index(UNK_WORD)
assert self.reloaded_params.mask_index == self.dico.index(MASK_WORD)
# build model / reload weights
self.reloaded_params['reload_model'] = ','.join([params.model_path] * 2)
encoder, decoder = build_model(self.reloaded_params, self.dico)
self.encoder = encoder[0]
self.encoder.load_state_dict(reloaded['encoder'])
assert len(reloaded['encoder'].keys()) == len(
list(p for p, _ in self.encoder.state_dict().items()))
self.decoder = decoder[0]
self.decoder.load_state_dict(reloaded['decoder'])
assert len(reloaded['decoder'].keys()) == len(
list(p for p, _ in self.decoder.state_dict().items()))
#self.encoder.to('cpu') #cuda()
#self.decoder.to('cpu') #cuda()
self.encoder.cuda()
self.decoder.cuda()
self.encoder.eval()
self.decoder.eval()
self.bpe_model = fastBPE.fastBPE(os.path.abspath(params.BPE_path))
def __init__(self, src_lang, tgt_lang):
model_path = TranscoderClient.get_model_path(src_lang, tgt_lang)
reloaded = torch.load(model_path, map_location='cpu')
reloaded['encoder'] = {(k[len('module.'):] if k.startswith('module.') else k): v for k, v in
reloaded['encoder'].items()}
assert 'decoder' in reloaded or (
'decoder_0' in reloaded and 'decoder_1' in reloaded)
if 'decoder' in reloaded:
decoders_names = ['decoder']
else:
decoders_names = ['decoder_0', 'decoder_1']
for decoder_name in decoders_names:
reloaded[decoder_name] = {(k[len('module.'):] if k.startswith('module.') else k): v for k, v in
reloaded[decoder_name].items()}
self.reloaded_params = AttrDict(reloaded['params'])
# build dictionary / update parameters
self.dico = Dictionary(
reloaded['dico_id2word'], reloaded['dico_word2id'], reloaded['dico_counts'])
assert self.reloaded_params.n_words == len(self.dico)
assert self.reloaded_params.bos_index == self.dico.index(BOS_WORD)
assert self.reloaded_params.eos_index == self.dico.index(EOS_WORD)
assert self.reloaded_params.pad_index == self.dico.index(PAD_WORD)
assert self.reloaded_params.unk_index == self.dico.index(UNK_WORD)
assert self.reloaded_params.mask_index == self.dico.index(MASK_WORD)
# build model / reload weights
self.reloaded_params['reload_model'] = ','.join([model_path] * 2)
encoder, decoder = build_model(self.reloaded_params, self.dico)
self.encoder = encoder[0]
self.encoder.load_state_dict(reloaded['encoder'])
assert len(reloaded['encoder'].keys()) == len(
list(p for p, _ in self.encoder.state_dict().items()))
self.decoder = decoder[0]
self.decoder.load_state_dict(reloaded['decoder'])
assert len(reloaded['decoder'].keys()) == len(
list(p for p, _ in self.decoder.state_dict().items()))
self.encoder.cuda()
self.decoder.cuda()
self.encoder.eval()
self.decoder.eval()
self.bpe_model = fastBPE.fastBPE(os.path.abspath(BPE_PATH))
self.allowed_languages = [lang.value for lang in Languages]