(self, ptm_model_path, device = "cpu")
| 9 | |
| 10 | class VectorizeModel: |
| 11 | def __init__(self, ptm_model_path, device = "cpu") -> None: |
| 12 | self.tokenizer = BertTokenizer.from_pretrained(ptm_model_path) |
| 13 | self.model = SimcseModel(pretrained_bert_path=ptm_model_path, pooling="cls") |
| 14 | self.model.eval() |
| 15 | |
| 16 | # self.DEVICE = torch.device('cuda' if torch.cuda.is_available() else "cpu") |
| 17 | self.DEVICE = device |
| 18 | logger.info(device) |
| 19 | self.model.to(self.DEVICE) |
| 20 | |
| 21 | self.pdist = nn.PairwiseDistance(2) |
| 22 | |
| 23 | def predict_vec(self,query): |
| 24 | q_id = self.tokenizer(query, max_length = 200, truncation=True, padding="max_length", return_tensors='pt') |
nothing calls this directly
no test coverage detected