| 42 | ClassificationDataset, ClassificationCollator, FastTextCollator,FastText, TextCNN, TextRNN, TextRCNN, DRNN, TextVDCNN, Transformer, DPCNN, AttentiveConvNet, RegionEmbedding |
| 43 | |
| 44 | class Predictor(object): |
| 45 | def __init__(self, config): |
| 46 | self.config = config |
| 47 | self.model_name = config.model_name |
| 48 | self.use_cuda = config.device.startswith("cuda") |
| 49 | self.dataset_name = "ClassificationDataset" |
| 50 | self.collate_name = "FastTextCollator" if self.model_name == "FastText" \ |
| 51 | else "ClassificationCollator" |
| 52 | self.dataset = globals()[self.dataset_name](config, [], mode="infer") |
| 53 | self.collate_fn = globals()[self.collate_name](config, len(self.dataset.label_map)) |
| 54 | self.model = Predictor._get_classification_model(self.model_name, self.dataset, config) |
| 55 | Predictor._load_checkpoint(config.eval.model_dir, self.model, self.use_cuda) |
| 56 | self.model.eval() |
| 57 | |
| 58 | @staticmethod |
| 59 | def _get_classification_model(model_name, dataset, conf): |
| 60 | model = globals()[model_name](dataset, conf) |
| 61 | model = model.cuda(conf.device) if conf.device.startswith("cuda") else model |
| 62 | return model |
| 63 | |
| 64 | @staticmethod |
| 65 | def _load_checkpoint(file_name, model, use_cuda): |
| 66 | if use_cuda: |
| 67 | checkpoint = torch.load(file_name, weights_only=True) |
| 68 | else: |
| 69 | checkpoint = torch.load(file_name, weights_only=True, map_location=lambda storage, loc: storage) |
| 70 | model.load_state_dict(checkpoint["state_dict"]) |
| 71 | |
| 72 | def predict(self, texts): |
| 73 | """ |
| 74 | input texts should be json objects |
| 75 | """ |
| 76 | with torch.no_grad(): |
| 77 | batch_texts = [self.dataset._get_vocab_id_list(json.loads(text)) for text in texts] |
| 78 | batch_texts = self.collate_fn(batch_texts) |
| 79 | logits = self.model(batch_texts) |
| 80 | if self.config.task_info.label_type != ClassificationType.MULTI_LABEL: |
| 81 | probs = torch.softmax(logits, dim=1) |
| 82 | else: |
| 83 | probs = torch.sigmoid(logits) |
| 84 | probs = probs.cpu().tolist() |
| 85 | return np.array(probs) |
| 86 | |
| 87 | if __name__ == "__main__": |
| 88 | config = Config(config_file=sys.argv[1]) |