INNER CODE UNIT · Python

Predict

JackHCC/Chinese-Text-Classification-PyTorch · pretrain_predict.py:21

class Predict:
    def __init__(self, model_name='bert', dataset='THUCNews'):
        self.x = import_module('models.' + model_name)
        self.config = self.x.Config(dataset)
        self.model = self.x.Model(self.config).to('cpu')
        self.model.load_state_dict(torch.load(self.config.save_path, map_location='cpu'))

    def build_predict_text(self, text):
        token = self.config.tokenizer.tokenize(text)
        token = ['[CLS]'] + token
        seq_len = len(token)
        mask = []
        token_ids = self.config.tokenizer.convert_tokens_to_ids(token)
        pad_size = self.config.pad_size
        if pad_size:
            if len(token) < pad_size:
                mask = [1] * len(token_ids) + ([0] * (pad_size - len(token)))
                token_ids += ([0] * (pad_size - len(token)))

View source record →