generate.py 文件源码

def __init__(self, model=None, model_source=None, src_dict=None, args=None):
        assert model is not None or model_source is not None

        if model is None:
            model_source = torch.load(model_source, map_location=lambda storage, loc: storage)
            self.dict = model_source["src_dict"]
            self.args = model_source["settings"]
            model = Model(self.args)
            model.load_state_dict(model_source['model'])
        else:
            self.dict = src_dict
            self.args = args

        self.num_directions = 2 if self.args.bidirectional else 1
        self.idx2word = {v: k for k, v in self.dict.items()}
        self.model = model.eval()