Vbert

colpali_engine/__init__.py

@@ -1,18 +1,22 @@
 from .models import (
     BiPali,
     BiPaliProj,
     BiQwen2,
     BiQwen2_5,
     BiQwen2_5_Processor,
     BiQwen2Processor,
+    BiModernVBert,
+    BiModernVBertProcessor,
     ColIdefics3,
     ColIdefics3Processor,
     ColPali,
     ColPaliProcessor,
     ColQwen2,
     ColQwen2_5,
     ColQwen2_5_Processor,
-    ColQwen2_5Omni,
-    ColQwen2_5OmniProcessor,
+    # ColQwen2_5Omni,
+    # ColQwen2_5OmniProcessor,
     ColQwen2Processor,
+    ColModernVBert,
+    ColModernVBertProcessor,
 )

colpali_engine/collators/visual_retriever_collator.py

@@ -1,11 +1,12 @@
 import random
 from typing import Any, Dict, List, Union
+import torch
 
 from PIL.Image import Image
 
 from colpali_engine.data.dataset import ColPaliEngineDataset
 from colpali_engine.models.paligemma import ColPaliProcessor
 from colpali_engine.utils.processing_utils import BaseVisualRetrieverProcessor
 
 
 def prefix_keys(data: Dict[str, Any], prefix: str) -> Dict[str, Any]:
@@ -69,16 +70,18 @@
 
             neg_tgt = example.get(ColPaliEngineDataset.NEG_TARGET_KEY, None)
             if neg_tgt is not None:
-                sampled_neg = random.choice(neg_tgt) if isinstance(neg_tgt, list) else neg_tgt
-                neg_targets.append(sampled_neg)
+                neg_targets.append(neg_tgt)
 
         # Ensure all queries are strings or images.
         assert all(isinstance(q, str) for q in queries), (
             "All queries must be strings, this collator does not support images in queries."
         )
 
+        is_str = isinstance(queries[0], str)
+
         # Process queries.
-        queries = [self.processor.query_prefix + q + self.processor.query_augmentation_token * 10 for q in queries]
+        # queries = [self.processor.query_prefix + q + self.processor.query_augmentation_token * 10 for q in queries]
+        queries = [q + self.processor.query_augmentation_token * 10 for q in queries] if is_str else queries
         batch_query = self.auto_collate(queries, key_prefix=self.query_prefix)
 
         # Process targets.
@@ -102,6 +105,26 @@
             proc_batch = self.processor.process_texts(texts=batch)
         elif isinstance(batch[0], Image):
             proc_batch = self.processor.process_images(images=batch)
+        elif isinstance(batch[0], list):
+            if isinstance(batch[0][0], str):
+                proc_texts_batch = []
+                batch_size = len(batch)
+                all_texts = [text for texts in batch for text in texts]
+                num_negatives = len(all_texts) // batch_size
+                proc_batch = self.processor.process_texts(texts=all_texts)
+            elif isinstance(batch[0][0], Image):
+                proc_imgs_batch = []
+                batch_size = len(batch)
+                all_imgs = [img for imgs in batch for img in imgs]
+                num_negatives = len(all_imgs) // batch_size
+                proc_batch = self.processor.process_images(images=all_imgs)
+            else:
+                raise ValueError(f"Unsupported batch type: {type(batch[0][0])}. Expected str or Image.")
+            for k, v in proc_batch.items():
+                if isinstance(v, torch.Tensor):                        
+                    proc_batch[k] = v.view(batch_size, num_negatives, *v.shape[1:])
+                else:
+                    proc_batch[k] = v
         else:
             raise ValueError(f"Unsupported batch type: {type(batch[0])}. Expected str or Image.")
         return prefix_keys(proc_batch, key_prefix)

colpali_engine/data/dataset.py

@@ -77,6 +77,7 @@ def __init__(
         query_column_name: str = "query",
         pos_target_column_name: str = "pos_target",
         neg_target_column_name: str = None,
+        num_negatives: int = 3,
     ):
         """
         Initialize the dataset with the provided data and external document corpus.
@@ -94,6 +95,7 @@ def __init__(
         self.pos_target_column_name = pos_target_column_name
         self.neg_target_column_name = neg_target_column_name
 
+        self.num_negatives = num_negatives
         assert isinstance(
             self.data,
             (list, Dataset, HFDataset),
@@ -131,8 +133,8 @@ def __getitem__(self, idx: int) -> Dict[str, Any]:
             pos_targets = [self.corpus.retrieve(doc_id) for doc_id in pos_targets]
             if neg_targets is not None:
                 # to avoid oveflowing CPU memory
-                if len(neg_targets) > 5:
-                    neg_targets = random.sample(neg_targets, 5)
+                if len(neg_targets) > self.num_negatives:
+                    neg_targets = random.sample(neg_targets, self.num_negatives)
                 neg_targets = [self.corpus.retrieve(doc_id) for doc_id in neg_targets]
 
         return {

colpali_engine/loss/__init__.py

@@ -4,11 +4,13 @@
     BiNegativeCELoss,
     BiPairwiseCELoss,
     BiPairwiseNegativeCELoss,
+    BiSigmoidLoss,
 )
 from .late_interaction_losses import (
     ColbertLoss,
     ColbertModule,
     ColbertNegativeCELoss,
     ColbertPairwiseCELoss,
     ColbertPairwiseNegativeCELoss,
+    ColbertSigmoidLoss,
 )

colpali_engine/loss/bi_encoder_losses.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn.functional as F  # noqa: N812
 from torch.nn import CrossEntropyLoss
 
 
@@ -111,6 +112,60 @@ def forward(
 
         return self.ce_loss(scores / self.temperature, pos_idx)
 
+class BiPairedEncoderLoss(BiEncoderModule):
+    """
+    InfoNCE loss for bi-encoders without explicit negatives.
+
+    Args:
+        temperature (float): Scaling factor for logits.
+        pos_aware_negative_filtering (bool): Apply in-batch negative filtering if True.
+        max_batch_size (int): Max batch size for index buffer caching.
+        filter_threshold (float): Threshold ratio for negative filtering.
+        filter_factor (float): Factor to down-weight filtered negatives.
+    """
+
+    def __init__(
+        self,
+        temperature: float = 0.02,
+        pos_aware_negative_filtering: bool = False,
+        max_batch_size: int = 1024,
+        filter_threshold: float = 0.95,
+        filter_factor: float = 0.5,
+    ):
+        super().__init__(max_batch_size, temperature, filter_threshold, filter_factor)
+        self.pos_aware_negative_filtering = pos_aware_negative_filtering
+        self.ce_loss = CrossEntropyLoss()
+
+    def forward(
+        self,
+        query_embeddings: torch.Tensor,
+        doc_embeddings: torch.Tensor,
+        offset: int = 0,
+    ) -> torch.Tensor:
+        """
+        Compute the InfoNCE loss over a batch of bi-encoder embeddings.
+
+        Args:
+            query_embeddings (Tensor[B, D]): Query vectors.
+            doc_embeddings (Tensor[B, D]): Document vectors.
+            offset (int): Offset for positive indices (multi-GPU).
+
+        Returns:
+            Tensor: Scalar cross-entropy loss.
+        """
+        # Compute in-batch similarity matrix
+        scores = torch.einsum("bd,cd->bc", query_embeddings, doc_embeddings)
+        batch_size = scores.size(0)
+        idx, pos_idx = self._get_idx(batch_size, offset, scores.device)
+
+        if self.pos_aware_negative_filtering:
+            self._filter_high_negatives(scores, pos_idx)
+
+        q2t = self.ce_loss(scores / self.temperature, pos_idx)
+        t2q = self.ce_loss(scores.T / self.temperature, ...)
+
+        return (q2t + t2q) / 2.0
+
 
 class BiNegativeCELoss(BiEncoderModule):
     """
@@ -161,17 +216,18 @@ def forward(
         Args:
             query_embeddings (Tensor[B, D]): Query vectors.
             doc_embeddings (Tensor[B, D]): Positive document vectors.
-            neg_doc_embeddings (Tensor[B, D]): Negative document vectors.
+            neg_doc_embeddings (Tensor[B, N, D]): Negative document vectors.
             offset (int): Offset for in-batch CE positives.
 
         Returns:
             Tensor: Scalar loss value.
         """
         # Dot-product only for matching pairs
-        pos_scores = (query_embeddings * doc_embeddings).sum(dim=1) / self.temperature
-        neg_scores = (query_embeddings * neg_doc_embeddings).sum(dim=1) / self.temperature
+        pos_scores = (query_embeddings * doc_embeddings[offset:offset + neg_doc_embeddings.size(0)]).sum(dim=1)
+        pos_scores /= self.temperature
+        neg_scores = torch.einsum("bd,bnd->bn", query_embeddings, neg_doc_embeddings) / self.temperature
 
-        loss = torch.nn.functional.softplus(neg_scores - pos_scores).mean()
+        loss = F.softplus(neg_scores - pos_scores.unsqueeze(1)).mean()
 
         if self.in_batch_term_weight > 0:
             loss_ib = self.inner_loss(query_embeddings, doc_embeddings, offset)
@@ -206,6 +262,7 @@ def forward(
         self,
         query_embeddings: torch.Tensor,
         doc_embeddings: torch.Tensor,
+        offset: int = 0,
     ) -> torch.Tensor:
         """
         Compute softplus(hardest_neg - pos) where hardest_neg is the highest off-diagonal score.
@@ -267,26 +324,93 @@ def forward(
         query_embeddings: torch.Tensor,
         doc_embeddings: torch.Tensor,
         neg_doc_embeddings: torch.Tensor,
+        offset: int = 0,
     ) -> torch.Tensor:
         """
         Compute softplus(neg-explicit - pos) plus optional pairwise in-batch loss.
 
         Args:
             query_embeddings (Tensor[B, D]): Query vectors.
             doc_embeddings (Tensor[B, D]): Positive document vectors.
-            neg_doc_embeddings (Tensor[B, D]): Negative document vectors.
+            neg_doc_embeddings (Tensor[B, N, D]): Negative document vectors.
 
         Returns:
             Tensor: Scalar loss value.
         """
         # dot product for matching pairs only
-        pos = (query_embeddings * doc_embeddings).sum(dim=1)
-        neg = (query_embeddings * neg_doc_embeddings).sum(dim=1)
+        pos = (query_embeddings * doc_embeddings).sum(dim=1) # B
+        neg = (query_embeddings.unsqueeze(1) * neg_doc_embeddings).sum(dim=2) # B x N
 
-        loss = torch.nn.functional.softplus((neg - pos) / self.temperature).mean()
+        loss = torch.nn.functional.softplus((neg - pos.unsqueeze(1)) / self.temperature).mean()
 
         if self.in_batch_term_weight > 0:
             loss_ib = self.inner_pairwise(query_embeddings, doc_embeddings)
             loss = loss * (1 - self.in_batch_term_weight) + loss_ib * self.in_batch_term_weight
 
         return loss
+
+class BiSigmoidLoss(BiEncoderModule):
+    """
+    Sigmoid loss for ColBERT with in-batch negatives.
+
+    Args:
+        temperature (float): Scaling factor for logits.
+        pos_aware_negative_filtering (bool): Apply in-batch negative filtering if True.
+        max_batch_size (int): Max batch size for index buffer caching.
+        filter_threshold (float): Threshold ratio for negative filtering.
+        filter_factor (float): Factor to down-weight filtered negatives.
+    """
+
+    def __init__(
+        self,
+        temperature: float = 0.02,
+        pos_aware_negative_filtering: bool = False,
+        max_batch_size: int = 1024,
+        filter_threshold: float = 0.95,
+        filter_factor: float = 0.5,
+    ):
+        super().__init__(max_batch_size, temperature, filter_threshold, filter_factor)
+        self.pos_aware_negative_filtering = pos_aware_negative_filtering
+
+    def forward(self, query_embeddings: torch.Tensor, doc_embeddings: torch.Tensor, offset: int = 0) -> torch.Tensor:
+        """
+        Compute the sigmoid loss for a batch of bi-encoder embeddings.
+
+        Args:
+            query_embeddings (Tensor[B, D]): Query vectors.
+            doc_embeddings (Tensor[B, D]): Document vectors.
+            offset (int): Offset for positive indices (multi-GPU).
+
+        Returns:
+            Tensor: Scalar cross-entropy loss.
+        """
+
+        # Compute in-batch similarity matrix
+        scores = torch.einsum("bd,cd->bc", query_embeddings, doc_embeddings)
+
+        batch_size, num_targets = scores.shape
+        device = scores.device
+
+        _, pos_idx = self._get_idx(batch_size, offset, device)
+
+        if self.pos_aware_negative_filtering:
+            self._filter_high_negatives(scores, pos_idx)
+
+        all_losses = []
+        for k in range(num_targets // batch_size):
+            # mask equal to 1 on offset -> offset + batch_size
+            curr_idx = torch.arange(offset, offset + batch_size, device=device)
+            # keep only the scores for the current batch
+            curr_scores = scores[:, curr_idx].view(-1) / self.temperature
+            # compute the labels
+            labels = -torch.ones(batch_size * batch_size, device=device)
+            if k == 0:
+                flat_pos = (pos_idx - offset) * (batch_size + 1)
+                labels[flat_pos] = 1.0
+            # compute the loss
+            block_loss = F.softplus(curr_scores * labels)
+            all_losses.append(block_loss)
+            # shift the offset for the next batch
+            offset = (offset + batch_size) % num_targets
+
+        return torch.stack(all_losses, dim=0).mean()