Merge pull request #161 from webis-de/docu

Enhance documentation for dataset classes and model configurations

Author: RaykKretzschmar

docs/guide/datasets.rst

@@ -4,8 +4,12 @@
 Which Dataset Format to Use?
 ===========================
 
-Lightning IR provides four dataset classes. The right one depends on your
-workflow and the shape of your data.
+Lightning IR provides four dataset classes —
+:py:class:`~lightning_ir.data.dataset.TupleDataset`,
+:py:class:`~lightning_ir.data.dataset.RunDataset`,
+:py:class:`~lightning_ir.data.dataset.DocDataset`, and
+:py:class:`~lightning_ir.data.dataset.QueryDataset` —
+which one to use depends on your workflow and the shape of your data.
 
 .. code-block:: text
 
@@ -49,15 +53,15 @@ Dataset Class Reference
    * - :py:class:`~lightning_ir.data.dataset.RunDataset`
      - ``fit``, ``re_rank``
      - Loads a ranked list of documents per query from a TREC-format run file
-       or an ir_datasets ID. Key parameters: ``depth`` (max rank to load),
+       or an `ir_datasets <https://ir-datasets.com/>`_ ID. Key parameters: ``depth`` (max rank to load),
        ``sample_size`` (docs per query), ``sampling_strategy``
        (``top`` or ``random``), ``targets`` (``relevance`` or ``score``).
    * - :py:class:`~lightning_ir.data.dataset.DocDataset`
      - ``index``
-     - Iterates over all documents in a collection. Backed by ir_datasets.
+     - Iterates over all documents in a collection. Backed by `ir_datasets <https://ir-datasets.com/>`_.
    * - :py:class:`~lightning_ir.data.dataset.QueryDataset`
      - ``search``
-     - Iterates over queries in a dataset split. Backed by ir_datasets.
+     - Iterates over queries in a dataset split. Backed by `ir_datasets <https://ir-datasets.com/>`_.
 
 .. tip::
 

docs/guide/index-types.rst

@@ -104,6 +104,9 @@ Index Type Comparison
 Quick Examples
 --------------
 
+Index Configs
+^^^^^^^^^^^^^
+
 **FAISS IVF index** (approximate nearest-neighbor for large dense collections):
 
 .. code-block:: yaml
@@ -202,6 +205,9 @@ Quick Examples
    )
    trainer.index(module, data_module)
 
+Search Configs
+^^^^^^^^^^^^^^
+
 **FAISS search** (querying a FAISS IVF index with a dense bi-encoder):
 
 .. code-block:: yaml

docs/guide/models.rst

@@ -79,19 +79,47 @@ Architecture Comparison
      - ❌
      - ✅
      - Sees all candidates at once; highest re-rank quality
+   * - **MVR**
+     - :py:class:`~lightning_ir.models.bi_encoders.mvr.MvrConfig`
+     - Separate
+     - Multi dense (viewer tokens)
+     - ✅
+     - ✅
+     - Fixed-count multi-vector; balances index size and quality
+   * - **COIL**
+     - :py:class:`~lightning_ir.models.bi_encoders.coil.CoilConfig`
+     - Separate
+     - Multi (token + CLS)
+     - ✅
+     - ✅
+     - Exact lexical match with context; sparse + dense hybrid
+   * - **UniCOIL**
+     - :py:class:`~lightning_ir.models.bi_encoders.coil.UniCoilConfig`
+     - Separate
+     - Single sparse (token weights)
+     - ✅
+     - ✅
+     - Lightweight token-weight sparse retrieval; simpler than SPLADE
 
 .. note::
 
-   Lightning IR also supports several other bi-encoder variants (:py:class:`~lightning_ir.models.bi_encoders.coil.CoilConfig`,
-   :py:class:`~lightning_ir.models.bi_encoders.mvr.MvrConfig`) and external models such as XTR for advanced use cases. See
-   the :ref:`concepts-model` page and the :py:mod:`~lightning_ir.models` API
-   reference for a full list.
+   External checkpoints are also supported out of the box. For example, XTR
+   (``google/xtr-base-en``) uses the ColBERT architecture (:py:class:`~lightning_ir.models.bi_encoders.col.ColConfig`)
+   with a T5 backbone. See the :ref:`concepts-model` page and the :py:mod:`~lightning_ir.models` API
+   reference for a full list of registered checkpoints.
 
 Quick Examples
 --------------
 
 **DPR bi-encoder** — simplest dense retrieval:
 
+DPR encodes each query and document into a single dense vector using a
+pooling step (typically the CLS token). Similarity is computed with a single
+dot product or cosine comparison, which makes retrieval fast and straightforward
+to index. The optional ``projection`` head lets you reduce the embedding
+dimension and add a non-linear bottleneck; set ``embedding_dim`` to control the
+output size.
+
 .. code-block:: yaml
 
    # model-dpr.yaml
@@ -101,6 +129,13 @@ Quick Examples
        model_name_or_path: bert-base-uncased
        config:
          class_path: lightning_ir.models.DprConfig
+         init_args:
+           similarity_function: dot   # or "cosine"
+           query_length: 32           # max query tokens
+           doc_length: 512            # max document tokens
+           pooling_strategy: first    # CLS token; also "mean", "max", "sum"
+           embedding_dim: null        # null → use backbone hidden size
+           projection: linear         # linear projection head; null to disable
 
 .. code-block:: python
 
@@ -109,11 +144,23 @@ Quick Examples
 
    module = BiEncoderModule(
        model_name_or_path="bert-base-uncased",
-       config=DprConfig(),
+       config=DprConfig(
+           similarity_function="dot",
+           query_length=32,
+           doc_length=512,
+           pooling_strategy="first",
+           embedding_dim=None,   # None → use backbone hidden size
+           projection="linear",
+       ),
    )
 
 **ColBERT** — multi-vector late interaction:
 
+ColBERT produces one embedding **per token**, so controlling sequence lengths and
+the per-token embedding dimension directly affects index size and quality.
+It also supports special ColBERT-specific options (query expansion, marker
+tokens, per-token ``normalization_strategy``) that DPR and SPLADE do not have.
+
 .. code-block:: yaml
 
    # model-colbert.yaml
@@ -126,13 +173,13 @@ Quick Examples
          init_args:
            similarity_function: dot
            query_aggregation_function: sum
-           query_expansion: true
+           query_expansion: true         # ColBERT-specific: pad queries to query_length
            query_length: 32
-           doc_length: 256
-           normalization_strategy: l2
-           embedding_dim: 128
-           projection: linear_no_bias
-           add_marker_tokens: true
+           doc_length: 256               # kept smaller than DPR to limit index size
+           normalization_strategy: l2   # ColBERT-specific: per-token l2 normalisation
+           embedding_dim: 128            # project every token to 128-d (reduces index size)
+           projection: linear_no_bias   # ColBERT convention: no bias in projection
+           add_marker_tokens: true       # ColBERT-specific: [Q]/[D] special tokens
 
 .. code-block:: python
 
@@ -144,18 +191,26 @@ Quick Examples
        config=ColConfig(
            similarity_function="dot",
            query_aggregation_function="sum",
-           query_expansion=True,
+           query_expansion=True,         # ColBERT-specific
            query_length=32,
            doc_length=256,
-           normalization_strategy="l2",
+           normalization_strategy="l2",  # ColBERT-specific
            embedding_dim=128,
            projection="linear_no_bias",
-           add_marker_tokens=True,
+           add_marker_tokens=True,       # ColBERT-specific
        ),
    )
 
 **SPLADE** — learned sparse retrieval:
 
+SPLADE maps each query and document to a sparse vector over the full vocabulary.
+Each vocabulary dimension is activated by taking the max-pool of the token
+logits, so the representation is directly interpretable as a bag of weighted
+terms. Because the embedding space is the tokenizer vocabulary, ``embedding_dim``
+cannot be set freely — it is always equal to the vocabulary size. The key
+knob is ``pooling_strategy``: ``max`` (the default) gives standard SPLADE
+behaviour; ``sum`` gives SPLADE-doc behaviour.
+
 .. code-block:: yaml
 
    # model-splade.yaml
@@ -165,6 +220,17 @@ Quick Examples
        model_name_or_path: bert-base-uncased
        config:
          class_path: lightning_ir.models.SpladeConfig
+         init_args:
+           similarity_function: dot   # sparse dot product
+           query_length: 32           # max query tokens
+           doc_length: 512            # max document tokens
+           pooling_strategy: max      # max over token activations (SPLADE default)
+
+.. note::
+
+   SPLADE uses the full vocabulary as its embedding space, so
+   ``embedding_dim`` is tied to the vocabulary size and cannot be configured
+   via the constructor (it is derived from the backbone tokenizer).
 
 .. code-block:: python
 
@@ -173,11 +239,23 @@ Quick Examples
 
    module = BiEncoderModule(
        model_name_or_path="bert-base-uncased",
-       config=SpladeConfig(),
+       config=SpladeConfig(
+           similarity_function="dot",
+           query_length=32,
+           doc_length=512,
+           pooling_strategy="max",  # max over token activations (SPLADE default)
+       ),
    )
 
 **Cross-encoder (MonoEncoder)** — highest quality re-ranking:
 
+A cross-encoder concatenates query and document into a single input and runs
+them jointly through the backbone, so every layer can attend across both texts.
+This yields the highest scoring quality but means no pre-indexing is possible
+— documents must be re-encoded for every new query. Use this architecture when
+you already have a candidate list (e.g. from a bi-encoder first stage) and
+need the best possible re-ranking without latency constraints.
+
 .. code-block:: yaml
 
    # model-cross-encoder.yaml