Make cross-encoder abstract (#76)

* fix: enhance documentation for SetEncoder config, model and tokenizer

* fix: enhance documentation for T5 cross-encoder configuration, model and tokenizer

* fix: update documentation in SetEncoderConfig for clarity on document embeddings

* feat: implement MonoModel for mono cross-encoder and update T5CrossEncoderModel to inherit from MonoModel

* fix pretty printing results for run files

* flake8 + black + add pylate dependency

* fix flake8

* iosrt

* fix dependencies in test ci

* add option to pass model kwargs to modules

* fix pretty printing for skipped inference datasets

* update config with model kwargs when necessary

* remove quotes when writing run

* register monoelectra rank-distillm data

* remove flash

* make model registration public

* update seismic test index

* remove test indexes

* Add MonoConfig and MonoModel classes; update model registration and tests

* finish making cross-encoder abstract

---------

Co-authored-by: Rayk Kretzschmar <112646288+RaykKretzschmar@users.noreply.github.com>

Author: fschlatt

lightning_ir/cross_encoder/cross_encoder_model.py

@@ -4,6 +4,7 @@
 This module defines the model class used to implement cross-encoder models.
 """
 
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Type
 
@@ -23,7 +24,7 @@ class CrossEncoderOutput(LightningIROutput):
     """Joint query-document embeddings"""
 
 
-class CrossEncoderModel(LightningIRModel):
+class CrossEncoderModel(LightningIRModel, ABC):
     config_class: Type[CrossEncoderConfig] = CrossEncoderConfig
     """Configuration class for cross-encoder models."""
 
@@ -36,9 +37,9 @@ def __init__(self, config: CrossEncoderConfig, *args, **kwargs):
         """
         super().__init__(config, *args, **kwargs)
         self.config: CrossEncoderConfig
-        self.linear = torch.nn.Linear(config.hidden_size, 1, bias=config.linear_bias)
 
     @batch_encoding_wrapper
+    @abstractmethod
     def forward(self, encoding: BatchEncoding) -> CrossEncoderOutput:
         """Computes contextualized embeddings for the joint query-document input sequence and computes a relevance
         score.
@@ -48,9 +49,4 @@ def forward(self, encoding: BatchEncoding) -> CrossEncoderOutput:
         :return: Output of the model
         :rtype: CrossEncoderOutput
         """
-        embeddings = self._backbone_forward(**encoding).last_hidden_state
-        embeddings = self.pooling(
-            embeddings, encoding.get("attention_mask", None), pooling_strategy=self.config.pooling_strategy
-        )
-        scores = self.linear(embeddings).view(-1)
-        return CrossEncoderOutput(scores=scores, embeddings=embeddings)
+        pass

lightning_ir/cross_encoder/cross_encoder_tokenizer.py

@@ -17,7 +17,9 @@ class CrossEncoderTokenizer(LightningIRTokenizer):
     config_class: Type[CrossEncoderConfig] = CrossEncoderConfig
     """Configuration class for the tokenizer."""
 
-    def __init__(self, *args, query_length: int = 32, doc_length: int = 512, **kwargs):
+    def __init__(
+        self, *args, query_length: int = 32, doc_length: int = 512, tokenizer_pattern: str | None = None, **kwargs
+    ):
         """:class:`.LightningIRTokenizer` for cross-encoder models. Encodes queries and documents jointly and ensures
         that the input sequences are of the correct length.
 
@@ -27,7 +29,10 @@ def __init__(self, *args, query_length: int = 32, doc_length: int = 512, **kwarg
         :type doc_length: int, optional
         :type doc_length: int, optional
         """
-        super().__init__(*args, query_length=query_length, doc_length=doc_length, **kwargs)
+        super().__init__(
+            *args, query_length=query_length, doc_length=doc_length, tokenizer_pattern=tokenizer_pattern, **kwargs
+        )
+        self.tokenizer_pattern = tokenizer_pattern
 
     def _truncate(self, text: Sequence[str], max_length: int) -> List[str]:
         """Encodes a list of texts, truncates them to a maximum number of tokens and decodes them to strings."""
@@ -98,22 +103,17 @@ def tokenize(
             raise ValueError("Both queries and docs must be provided.")
         if isinstance(docs, str) and not isinstance(queries, str):
             raise ValueError("Queries and docs must be both lists or both strings.")
-        is_string_queries = False
-        is_string_docs = False
         if isinstance(queries, str):
             queries = [queries]
-            is_string_queries = True
         if isinstance(docs, str):
             docs = [docs]
-            is_string_docs = True
-        is_string_both = is_string_queries and is_string_docs
         num_docs = self._process_num_docs(queries, docs, num_docs)
         queries, docs = self._preprocess(queries, docs, num_docs)
-        return_tensors = kwargs.get("return_tensors", None)
-        if return_tensors is not None:
-            kwargs["pad_to_multiple_of"] = 8
-        if is_string_both:
-            encoding = self(queries[0], docs[0], **kwargs)
+
+        if self.tokenizer_pattern is not None:
+            input_texts = [self.tokenizer_pattern.format(query=query, doc=doc) for query, doc in zip(queries, docs)]
+            encoding = self(input_texts, **kwargs)
         else:
             encoding = self(queries, docs, **kwargs)
+
         return {"encoding": encoding}

lightning_ir/models/__init__.py

@@ -1,21 +1,20 @@
 from .col import ColConfig, ColModel, ColTokenizer
 from .dpr import DprConfig, DprModel
+from .mono import MonoConfig, MonoModel
 from .set_encoder import SetEncoderConfig, SetEncoderModel, SetEncoderTokenizer
 from .splade import SpladeConfig, SpladeModel
-from .t5_cross_encoder import T5CrossEncoderConfig, T5CrossEncoderModel, T5CrossEncoderTokenizer
 
 __all__ = [
     "ColConfig",
     "ColModel",
     "ColTokenizer",
     "DprConfig",
     "DprModel",
+    "MonoConfig",
+    "MonoModel",
     "SetEncoderConfig",
     "SetEncoderModel",
     "SetEncoderTokenizer",
     "SpladeConfig",
     "SpladeModel",
-    "T5CrossEncoderConfig",
-    "T5CrossEncoderModel",
-    "T5CrossEncoderTokenizer",
 ]

lightning_ir/models/mono.py

@@ -0,0 +1,117 @@
+"""
+Model implementation for mono cross-encoder models. Originally introduced in
+`Passage Re-ranking with BERT
+<https://arxiv.org/abs/1901.04085>`_.
+"""
+
+from typing import Literal, Type
+
+import torch
+from transformers import BatchEncoding
+
+from ..base.model import batch_encoding_wrapper
+from ..cross_encoder.cross_encoder_config import CrossEncoderConfig
+from ..cross_encoder.cross_encoder_model import CrossEncoderModel, CrossEncoderOutput
+
+
+class ScaleLinear(torch.nn.Linear):
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586 # noqa
+        input = input * (input.shape[-1] ** -0.5)
+        return super().forward(input)
+
+
+class MonoConfig(CrossEncoderConfig):
+    """Configuration class for mono cross-encoder models."""
+
+    model_type = "mono"
+    """Model type for mono cross-encoder models."""
+
+    def __init__(
+        self,
+        query_length: int = 32,
+        doc_length: int = 512,
+        pooling_strategy: Literal["first", "mean", "max", "sum", "bert_pool"] = "first",
+        linear_bias: bool = False,
+        scoring_strategy: Literal["mono", "rank"] = "rank",
+        tokenizer_pattern: str | None = None,
+        **kwargs,
+    ):
+        """Initialize the configuration for mono cross-encoder models."""
+        self._bert_pool = False
+        if pooling_strategy == "bert_pool":
+            self._bert_pool = True
+            pooling_strategy = "first"
+        super().__init__(
+            query_length=query_length,
+            doc_length=doc_length,
+            pooling_strategy=pooling_strategy,
+            linear_bias=linear_bias,
+            **kwargs,
+        )
+        self.scoring_strategy = scoring_strategy
+        self.tokenizer_pattern = tokenizer_pattern
+
+
+class MonoModel(CrossEncoderModel):
+    config_class: Type[MonoConfig] = MonoConfig
+    """Configuration class for mono cross-encoder models."""
+
+    def __init__(self, config: MonoConfig, *args, **kwargs):
+        """A cross-encoder model that jointly encodes a query and document(s). The contextualized embeddings are
+        aggragated into a single vector and fed to a linear layer which computes a final relevance score.
+
+        :param config: Configuration for the cross-encoder model
+        :type config: CrossEncoderConfig
+        """
+        super().__init__(config, *args, **kwargs)
+
+        if self.config.scoring_strategy == "mono":
+            output_dim = 2
+        elif self.config.scoring_strategy == "rank":
+            output_dim = 1
+        else:
+            raise ValueError(
+                f"Unknown scoring strategy {self.config.scoring_strategy}. Supported strategies are 'mono' and 'rank'."
+            )
+
+        self.bert_pool = torch.nn.Identity()
+        if self.config._bert_pool:
+            self.bert_pool = torch.nn.Sequential(
+                torch.nn.Linear(config.hidden_size, config.hidden_size), torch.nn.Tanh()
+            )
+
+        if self.config.backbone_model_type == "t5":
+            self.linear = ScaleLinear(config.hidden_size, output_dim, bias=self.config.linear_bias)
+        else:
+            self.linear = torch.nn.Linear(config.hidden_size, output_dim, bias=self.config.linear_bias)
+
+    @batch_encoding_wrapper
+    def forward(self, encoding: BatchEncoding) -> CrossEncoderOutput:
+        """Computes contextualized embeddings for the joint query-document input sequence and computes a relevance
+        score.
+
+        :param encoding: Tokenizer encoding for the joint query-document input sequence
+        :type encoding: BatchEncoding
+        :return: Output of the model
+        :rtype: CrossEncoderOutput
+        """
+        if hasattr(self, "decoder"):
+            # NOTE hack to make T5 cross-encoders work. other encoder-decoder models may not have `decoder` as their
+            # attribute. maybe find a better way to check for this?
+            decoder_input_ids = torch.zeros(
+                (encoding["input_ids"].shape[0], 1), device=encoding["input_ids"].device, dtype=torch.long
+            )
+            encoding["decoder_input_ids"] = decoder_input_ids
+        embeddings = self._backbone_forward(**encoding).last_hidden_state
+        embeddings = self.pooling(
+            embeddings, encoding.get("attention_mask", None), pooling_strategy=self.config.pooling_strategy
+        )
+        embeddings = self.bert_pool(embeddings)
+        scores = self.linear(embeddings)
+
+        if self.config.scoring_strategy == "mono":
+            scores = torch.nn.functional.log_softmax(scores.view(-1, 2), dim=-1)[:, 1]
+
+        return CrossEncoderOutput(scores=scores.view(-1), embeddings=embeddings)