predict.py 文件源码

def __init__(self, model_source, cuda=False, beam_size=3):
        self.torch = torch.cuda if cuda else torch
        self.cuda = cuda
        self.jb = Jieba("./segmenter_dicts", useSynonym=True, HMM=False)
        self.swf = StopwordFilter("./segmenter_dicts/stopwords.txt")

        model_source = torch.load(model_source)
        self.src_dict = model_source["src_dict"]
        self.tgt_dict = model_source["tgt_dict"]
        self.src_idx2ind = {v: k for k, v in model_source["tgt_dict"].items()}
        self.args = args = model_source["settings"]
        model = CNN_Ranking(args)
        model.load_state_dict(model_source['model'])

        if self.cuda:
            model = model.cuda()
        else:
            model = model.cpu()
        self.model = model.eval()