(self, query_list: Union[List[str], str], is_query=True)
| 36 | |
| 37 | @torch.inference_mode() |
| 38 | def single_batch_encode(self, query_list: Union[List[str], str], is_query=True) -> np.ndarray: |
| 39 | query_list = parse_query(self.model_name, query_list, self.instruction, is_query) |
| 40 | |
| 41 | inputs = self.tokenizer( |
| 42 | query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors="pt" |
| 43 | ) |
| 44 | inputs = {k: v.cuda() for k, v in inputs.items()} |
| 45 | |
| 46 | if "T5" in type(self.model).__name__ or (isinstance(self.model, torch.nn.DataParallel) and "T5" in type(self.model.module).__name__): |
| 47 | # T5-based retrieval model |
| 48 | decoder_input_ids = torch.zeros((inputs["input_ids"].shape[0], 1), dtype=torch.long).to( |
| 49 | inputs["input_ids"].device |
| 50 | ) |
| 51 | output = self.model(**inputs, decoder_input_ids=decoder_input_ids, return_dict=True) |
| 52 | query_emb = output.last_hidden_state[:, 0, :] |
| 53 | |
| 54 | else: |
| 55 | output = self.model(**inputs, return_dict=True) |
| 56 | pooler_output = output.get('pooler_output', None) |
| 57 | last_hidden_state = output.get('last_hidden_state', None) |
| 58 | query_emb = pooling( |
| 59 | pooler_output, last_hidden_state, inputs["attention_mask"], self.pooling_method |
| 60 | ) |
| 61 | if "dpr" not in self.model_name: |
| 62 | query_emb = torch.nn.functional.normalize(query_emb, dim=-1) |
| 63 | query_emb = query_emb.detach().cpu().numpy() |
| 64 | query_emb = query_emb.astype(np.float32, order="C") |
| 65 | return query_emb |
| 66 | |
| 67 | @torch.inference_mode() |
| 68 | def encode(self, query_list: List[str], batch_size=64, is_query=True) -> np.ndarray: |
no test coverage detected