| 8 | from src.models.vec_model.simcse_model import SimcseModel |
| 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') |
| 25 | with torch.no_grad(): |
| 26 | q_id_input_ids = q_id["input_ids"].squeeze(1).to(self.DEVICE) |
| 27 | q_id_attention_mask = q_id["attention_mask"].squeeze(1).to(self.DEVICE) |
| 28 | q_id_token_type_ids = q_id["token_type_ids"].squeeze(1).to(self.DEVICE) |
| 29 | q_id_pred = self.model(q_id_input_ids, q_id_attention_mask, q_id_token_type_ids) |
| 30 | |
| 31 | return q_id_pred |
| 32 | |
| 33 | def predict_vec_request(self, query): |
| 34 | q_id_pred = self.predict_vec(query) |
| 35 | return q_id_pred.cpu().numpy().tolist() |
| 36 | |
| 37 | def predict_sim(self, q1, q2): |
| 38 | q1_v = self.predict_vec(q1) |
| 39 | q2_v = self.predict_vec(q2) |
| 40 | sim = F.cosine_similarity(q1_v[0], q2_v[0], dim=-1) |
| 41 | return sim.numpy().tolist() |
| 42 | |
| 43 | if __name__ == "__main__": |
| 44 | import time,random |
no outgoing calls
no test coverage detected