diff --git a/FlagEmbedding/inference/embedder/encoder_only/m3.py b/FlagEmbedding/inference/embedder/encoder_only/m3.py index 2082003a..39fcb6d9 100644 --- a/FlagEmbedding/inference/embedder/encoder_only/m3.py +++ b/FlagEmbedding/inference/embedder/encoder_only/m3.py @@ -391,11 +391,12 @@ def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list): flag = False while flag is False: try: - inputs_batch = self.tokenizer.pad( - all_inputs_sorted[: batch_size], + inputs_batch = self.tokenizer( + sentences[start_index:start_index + batch_size], padding=True, + truncation=True, + max_length=max_length, return_tensors='pt', - **kwargs ).to(device) outputs = self.model( inputs_batch, @@ -413,10 +414,12 @@ def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list): all_dense_embeddings, all_lexical_weights, all_colbert_vecs = [], [], [] for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings", disable=len(sentences) < batch_size): - inputs_batch = all_inputs_sorted[start_index:start_index + batch_size] - inputs_batch = self.tokenizer.pad( - inputs_batch, + sentences_batch = [sentences[i] for i in length_sorted_idx[start_index:start_index + batch_size]] + inputs_batch = self.tokenizer( + sentences_batch, padding=True, + truncation=True, + max_length=max_length, return_tensors='pt', **kwargs ).to(device)