INNER CODE UNIT · Python

predict

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

    def predict(self, query):
        # 返回预测的索引
        data = self.build_predict_text(query)
        with torch.no_grad():
            outputs = self.model(data)
            num = torch.argmax(outputs)
        return key[int(num)]

    def predict_list(self, querys):
        pred = []
        for query in querys:
            pred.append(self.predict(query))
        return pred


if __name__ == "__main__":
    pred = Predict('bert')
    # 预测一条

View source record →