kaiko-eva 0.0.2__py3-none-any.whl → 0.1.1__py3-none-any.whl

eva/core/callbacks/__init__.py

@@ -1,6 +1,6 @@
 """Callbacks API."""
 
 from eva.core.callbacks.config import ConfigurationLogger
-from eva.core.callbacks.writers import EmbeddingsWriter
+from eva.core.callbacks.writers import ClassificationEmbeddingsWriter, SegmentationEmbeddingsWriter
 
-__all__ = ["ConfigurationLogger", "EmbeddingsWriter"]
+__all__ = ["ConfigurationLogger", "ClassificationEmbeddingsWriter", "SegmentationEmbeddingsWriter"]

eva/core/callbacks/writers/__init__.py

@@ -1,5 +1,8 @@
-"""Callbacks API."""
+"""Writers callbacks API."""
 
-from eva.core.callbacks.writers.embeddings import EmbeddingsWriter
+from eva.core.callbacks.writers.embeddings import (
+    ClassificationEmbeddingsWriter,
+    SegmentationEmbeddingsWriter,
+)
 
-__all__ = ["EmbeddingsWriter"]
+__all__ = ["ClassificationEmbeddingsWriter", "SegmentationEmbeddingsWriter"]

eva/core/callbacks/writers/embeddings/__init__.py

@@ -0,0 +1,6 @@
+"""Embedding callback writers."""
+
+from eva.core.callbacks.writers.embeddings.classification import ClassificationEmbeddingsWriter
+from eva.core.callbacks.writers.embeddings.segmentation import SegmentationEmbeddingsWriter
+
+__all__ = ["ClassificationEmbeddingsWriter", "SegmentationEmbeddingsWriter"]

eva/core/callbacks/writers/embeddings/_manifest.py

@@ -0,0 +1,71 @@
+"""Manifest file manager."""
+
+import csv
+import io
+import os
+from typing import Any, Dict, List
+
+import _csv
+import torch
+
+
+class ManifestManager:
+    """Class for writing the embedding manifest files."""
+
+    def __init__(
+        self,
+        output_dir: str,
+        metadata_keys: List[str] | None = None,
+        overwrite: bool = False,
+    ) -> None:
+        """Initializes the writing manager.
+
+        Args:
+            output_dir: The directory where the embeddings will be saved.
+            metadata_keys: An optional list of keys to extract from the batch
+                metadata and store as additional columns in the manifest file.
+            overwrite: Whether to overwrite the output directory.
+        """
+        self._output_dir = output_dir
+        self._metadata_keys = metadata_keys or []
+        self._overwrite = overwrite
+
+        self._manifest_file: io.TextIOWrapper
+        self._manifest_writer: _csv.Writer  # type: ignore
+
+        self._setup()
+
+    def _setup(self) -> None:
+        """Initializes the manifest file and sets the file object and writer."""
+        manifest_path = os.path.join(self._output_dir, "manifest.csv")
+        if os.path.exists(manifest_path) and not self._overwrite:
+            raise FileExistsError(
+                f"A manifest file already exists at {manifest_path}, which indicates that the "
+                "chosen output directory has been previously used for writing embeddings."
+            )
+        self._manifest_file = open(manifest_path, "w", newline="")
+        self._manifest_writer = csv.writer(self._manifest_file)
+        self._manifest_writer.writerow(
+            ["origin", "embeddings", "target", "split"] + self._metadata_keys
+        )
+
+    def update(
+        self,
+        input_name: str,
+        save_name: str,
+        target: str,
+        split: str | None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> None:
+        """Adds a new entry to the manifest file."""
+        metadata_entries = _to_dict_values(metadata or {})
+        self._manifest_writer.writerow([input_name, save_name, target, split] + metadata_entries)
+
+    def close(self) -> None:
+        """Closes the manifest file."""
+        if self._manifest_file:
+            self._manifest_file.close()
+
+
+def _to_dict_values(data: Dict[str, Any]) -> List[Any]:
+    return [value.item() if isinstance(value, torch.Tensor) else value for value in data.values()]

eva/core/callbacks/writers/embeddings/base.py

@@ -0,0 +1,192 @@
+"""Embeddings writer base class."""
+
+import abc
+import io
+import os
+from typing import Any, Dict, List, Sequence
+
+import lightning.pytorch as pl
+import torch
+from lightning.pytorch import callbacks
+from loguru import logger
+from torch import multiprocessing, nn
+from typing_extensions import override
+
+from eva.core import utils
+from eva.core.callbacks.writers.embeddings.typings import QUEUE_ITEM
+from eva.core.models.modules.typings import INPUT_BATCH
+from eva.core.utils import multiprocessing as eva_multiprocessing
+
+
+class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
+    """Callback for writing generated embeddings to disk."""
+
+    def __init__(
+        self,
+        output_dir: str,
+        backbone: nn.Module | None = None,
+        dataloader_idx_map: Dict[int, str] | None = None,
+        metadata_keys: List[str] | None = None,
+        overwrite: bool = False,
+        save_every_n: int = 100,
+    ) -> None:
+        """Initializes a new EmbeddingsWriter instance.
+
+        This callback writes the embedding files in a separate process to avoid blocking the
+        main process where the model forward pass is executed.
+
+        Args:
+            output_dir: The directory where the embeddings will be saved.
+            backbone: A model to be used as feature extractor. If `None`,
+                it will be expected that the input batch returns the features directly.
+            dataloader_idx_map: A dictionary mapping dataloader indices to their respective
+                names (e.g. train, val, test).
+            metadata_keys: An optional list of keys to extract from the batch metadata and store
+                as additional columns in the manifest file.
+            overwrite: Whether to overwrite if embeddings are already present in the specified
+                output directory. If set to `False`, an error will be raised if embeddings are
+                already present (recommended).
+            save_every_n: Interval for number of iterations to save the embeddings to disk.
+                During this interval, the embeddings are accumulated in memory.
+        """
+        super().__init__(write_interval="batch")
+
+        self._output_dir = output_dir
+        self._backbone = backbone
+        self._dataloader_idx_map = dataloader_idx_map or {}
+        self._overwrite = overwrite
+        self._save_every_n = save_every_n
+        self._metadata_keys = metadata_keys or []
+
+        self._write_queue: multiprocessing.Queue
+        self._write_process: eva_multiprocessing.Process
+
+    @staticmethod
+    @abc.abstractmethod
+    def _process_write_queue(
+        write_queue: multiprocessing.Queue,
+        output_dir: str,
+        metadata_keys: List[str],
+        save_every_n: int,
+        overwrite: bool = False,
+    ) -> None:
+        """This function receives and processes items added by the main process to the queue.
+
+        Queue items contain the embedding tensors, targets and metadata which need to be
+        saved to disk (.pt files and manifest).
+        """
+
+    @override
+    def on_predict_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        self._check_if_exists()
+        self._initialize_write_process()
+        self._write_process.start()
+
+        if self._backbone is not None:
+            self._backbone = self._backbone.to(pl_module.device)
+            self._backbone.eval()
+
+    @override
+    def write_on_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        prediction: Any,
+        batch_indices: Sequence[int],
+        batch: INPUT_BATCH,
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        dataset = trainer.predict_dataloaders[dataloader_idx].dataset  # type: ignore
+        _, targets, metadata = INPUT_BATCH(*batch)
+        split = self._dataloader_idx_map.get(dataloader_idx)
+        if not isinstance(targets, torch.Tensor):
+            raise ValueError(f"Targets ({type(targets)}) should be `torch.Tensor`.")
+
+        with torch.no_grad():
+            embeddings = self._get_embeddings(prediction)
+
+        for local_idx, global_idx in enumerate(batch_indices[: len(embeddings)]):
+            data_name = dataset.filename(global_idx)
+            save_name = os.path.splitext(data_name)[0] + ".pt"
+            embeddings_buffer, target_buffer = _as_io_buffers(
+                embeddings[local_idx], targets[local_idx]
+            )
+            item_metadata = self._get_item_metadata(metadata, local_idx)
+            item = QUEUE_ITEM(
+                prediction_buffer=embeddings_buffer,
+                target_buffer=target_buffer,
+                data_name=data_name,
+                save_name=save_name,
+                split=split,
+                metadata=item_metadata,
+            )
+            self._write_queue.put(item)
+
+        self._write_process.check_exceptions()
+
+    @override
+    def on_predict_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        self._write_queue.put(None)
+        self._write_process.join()
+        logger.info(f"Predictions and manifest saved to {self._output_dir}")
+
+    def _initialize_write_process(self) -> None:
+        self._write_queue = multiprocessing.Queue()
+        self._write_process = eva_multiprocessing.Process(
+            target=self._process_write_queue,
+            args=(
+                self._write_queue,
+                self._output_dir,
+                self._metadata_keys,
+                self._save_every_n,
+                self._overwrite,
+            ),
+        )
+
+    @abc.abstractmethod
+    def _get_embeddings(self, tensor: torch.Tensor) -> torch.Tensor | List[List[torch.Tensor]]:
+        """Returns the embeddings from predictions."""
+
+    def _get_item_metadata(
+        self, metadata: Dict[str, Any] | None, local_idx: int
+    ) -> Dict[str, Any] | None:
+        """Returns the metadata for the item at the given local index."""
+        if not metadata:
+            if self._metadata_keys:
+                raise ValueError("Metadata keys are provided but the batch metadata is empty.")
+            else:
+                return None
+
+        item_metadata = {}
+        for key in self._metadata_keys:
+            if key not in metadata:
+                raise KeyError(f"Metadata key '{key}' not found in the batch metadata.")
+            metadata_value = metadata[key][local_idx]
+            try:
+                item_metadata[key] = utils.to_cpu(metadata_value)
+            except TypeError:
+                item_metadata[key] = metadata_value
+
+        return item_metadata
+
+    def _check_if_exists(self) -> None:
+        """Checks if the output directory already exists and if it should be overwritten."""
+        try:
+            os.makedirs(self._output_dir, exist_ok=self._overwrite)
+        except FileExistsError as e:
+            raise FileExistsError(
+                f"The embeddings output directory already exists: {self._output_dir}. This "
+                "either means that they have been computed before or that a wrong output "
+                "directory is being used. Consider using `eva fit` instead, selecting a "
+                "different output directory or setting overwrite=True."
+            ) from e
+        os.makedirs(self._output_dir, exist_ok=True)
+
+
+def _as_io_buffers(*items: torch.Tensor | List[torch.Tensor]) -> Sequence[io.BytesIO]:
+    """Loads torch tensors as io buffers."""
+    buffers = [io.BytesIO() for _ in range(len(items))]
+    for tensor, buffer in zip(items, buffers, strict=False):
+        torch.save(utils.clone(tensor), buffer)
+    return buffers

eva/core/callbacks/writers/embeddings/classification.py

@@ -0,0 +1,117 @@
+"""Embeddings writer for classification."""
+
+import io
+import os
+from typing import Dict, List
+
+import torch
+from torch import multiprocessing
+from typing_extensions import override
+
+from eva.core.callbacks.writers.embeddings import base
+from eva.core.callbacks.writers.embeddings._manifest import ManifestManager
+from eva.core.callbacks.writers.embeddings.typings import ITEM_DICT_ENTRY, QUEUE_ITEM
+
+
+class ClassificationEmbeddingsWriter(base.EmbeddingsWriter):
+    """Callback for writing generated embeddings to disk for classification tasks."""
+
+    @staticmethod
+    @override
+    def _process_write_queue(
+        write_queue: multiprocessing.Queue,
+        output_dir: str,
+        metadata_keys: List[str],
+        save_every_n: int,
+        overwrite: bool = False,
+    ) -> None:
+        """Processes the write queue and saves the predictions to disk.
+
+        Note that in Multi Instance Learning (MIL) scenarios, we can have multiple
+        embeddings per input data point. In that case, this function will save all
+        embeddings that correspond to the same data point as a list of tensors to
+        the same .pt file.
+        """
+        manifest_manager = ManifestManager(output_dir, metadata_keys, overwrite)
+        name_to_items: Dict[str, ITEM_DICT_ENTRY] = {}
+
+        counter = 0
+        while True:
+            item = write_queue.get()
+            if item is None:
+                break
+            item = QUEUE_ITEM(*item)
+
+            if item.save_name in name_to_items:
+                name_to_items[item.save_name].items.append(item)
+            else:
+                name_to_items[item.save_name] = ITEM_DICT_ENTRY(items=[item], save_count=0)
+
+            if counter > 0 and counter % save_every_n == 0:
+                name_to_items = _save_items(name_to_items, output_dir, manifest_manager)
+            counter += 1
+
+        if len(name_to_items) > 0:
+            _save_items(name_to_items, output_dir, manifest_manager)
+
+        manifest_manager.close()
+
+    @override
+    def _get_embeddings(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Returns the embeddings from predictions."""
+        return self._backbone(tensor) if self._backbone else tensor
+
+
+def _save_items(
+    name_to_items: Dict[str, ITEM_DICT_ENTRY],
+    output_dir: str,
+    manifest_manager: ManifestManager,
+) -> Dict[str, ITEM_DICT_ENTRY]:
+    """Saves predictions to disk and updates the manifest file.
+
+    Args:
+        name_to_items: A dictionary mapping save data point names to the corresponding queue items
+            holding the prediction tensors and the information for the manifest file.
+        output_dir: The directory where the embedding tensors & manifest will be saved.
+        manifest_manager: The manifest manager instance to update the manifest file.
+    """
+    for save_name, entry in name_to_items.items():
+        if len(entry.items) > 0:
+            save_path = os.path.join(output_dir, save_name)
+            is_first_save = entry.save_count == 0
+            if is_first_save:
+                _, target, input_name, _, split, metadata = QUEUE_ITEM(*entry.items[0])
+                target = torch.load(io.BytesIO(target.getbuffer()), map_location="cpu").item()
+                manifest_manager.update(input_name, save_name, target, split, metadata)
+
+            prediction_buffers = [item.prediction_buffer for item in entry.items]
+            _save_predictions(prediction_buffers, save_path, is_first_save)
+            name_to_items[save_name].save_count += 1
+            name_to_items[save_name].items = []
+
+    return name_to_items
+
+
+def _save_predictions(
+    prediction_buffers: List[io.BytesIO], save_path: str, is_first_save: bool
+) -> None:
+    """Saves the embedding tensors as list to .pt files.
+
+    If it's not the first save to this save_path, the new predictions are appended to
+    the existing ones and saved to the same file.
+
+    Example use-case: Save all patch embeddings corresponding to the same WSI to a single file.
+    """
+    predictions = [
+        torch.load(io.BytesIO(buffer.getbuffer()), map_location="cpu")
+        for buffer in prediction_buffers
+    ]
+
+    if not is_first_save:
+        previous_predictions = torch.load(save_path, map_location="cpu")
+        if not isinstance(previous_predictions, list):
+            raise ValueError("Previous predictions should be a list of tensors.")
+        predictions = predictions + previous_predictions
+
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    torch.save(predictions, save_path)

eva/core/callbacks/writers/embeddings/segmentation.py

@@ -0,0 +1,78 @@
+"""Segmentation embeddings writer."""
+
+import collections
+import io
+import os
+from typing import List
+
+import torch
+from torch import multiprocessing
+from typing_extensions import override
+
+from eva.core.callbacks.writers.embeddings import base
+from eva.core.callbacks.writers.embeddings._manifest import ManifestManager
+from eva.core.callbacks.writers.embeddings.typings import QUEUE_ITEM
+
+
+class SegmentationEmbeddingsWriter(base.EmbeddingsWriter):
+    """Callback for writing generated embeddings to disk."""
+
+    @staticmethod
+    @override
+    def _process_write_queue(
+        write_queue: multiprocessing.Queue,
+        output_dir: str,
+        metadata_keys: List[str],
+        save_every_n: int,
+        overwrite: bool = False,
+    ) -> None:
+        manifest_manager = ManifestManager(output_dir, metadata_keys, overwrite)
+        counter = collections.defaultdict(lambda: -1)
+        while True:
+            item = write_queue.get()
+            if item is None:
+                break
+
+            embeddings_buffer, target_buffer, input_name, save_name, split, metadata = QUEUE_ITEM(
+                *item
+            )
+            counter[save_name] += 1
+            save_name = save_name.replace(".pt", f"-{counter[save_name]}.pt")
+            target_filename = save_name.replace(".pt", "-mask.pt")
+
+            _save_embedding(embeddings_buffer, save_name, output_dir)
+            _save_embedding(target_buffer, target_filename, output_dir)
+            manifest_manager.update(input_name, save_name, target_filename, split, metadata)
+
+        manifest_manager.close()
+
+    @override
+    def _get_embeddings(self, tensor: torch.Tensor) -> torch.Tensor | List[List[torch.Tensor]]:
+        """Returns the embeddings from predictions."""
+
+        def _get_grouped_embeddings(embeddings: List[torch.Tensor]) -> List[List[torch.Tensor]]:
+            """Casts a list of multi-leveled batched embeddings to grouped per batch.
+
+            That is, for embeddings to be a list of shape (batch_size, hidden_dim, height, width),
+            such as `[(2, 192, 16, 16), (2, 192, 16, 16)]`, to be reshaped as a list of lists of
+            per batch multi-level embeddings, thus
+            `[ [(192, 16, 16), (192, 16, 16)], [(192, 16, 16), (192, 16, 16)] ]`.
+            """
+            batch_size = embeddings[0].shape[0]
+            grouped_embeddings = []
+            for batch_idx in range(batch_size):
+                batch_list = [layer_embeddings[batch_idx] for layer_embeddings in embeddings]
+                grouped_embeddings.append(batch_list)
+            return grouped_embeddings
+
+        embeddings = self._backbone(tensor) if self._backbone else tensor
+        if isinstance(embeddings, list):
+            embeddings = _get_grouped_embeddings(embeddings)
+        return embeddings
+
+
+def _save_embedding(embeddings_buffer: io.BytesIO, save_name: str, output_dir: str) -> None:
+    save_path = os.path.join(output_dir, save_name)
+    prediction = torch.load(io.BytesIO(embeddings_buffer.getbuffer()), map_location="cpu")
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    torch.save(prediction, save_path)

eva/core/callbacks/writers/embeddings/typings.py

@@ -0,0 +1,38 @@
+"""Typing definitions for the writer callback functions."""
+
+import dataclasses
+import io
+from typing import Any, Dict, List, NamedTuple
+
+
+class QUEUE_ITEM(NamedTuple):
+    """The default input batch data scheme."""
+
+    prediction_buffer: io.BytesIO
+    """IO buffer containing the prediction tensor."""
+
+    target_buffer: io.BytesIO
+    """IO buffer containing the target tensor."""
+
+    data_name: str
+    """Name of the input data that was used to generate the embedding."""
+
+    save_name: str
+    """Name to store the generated embedding."""
+
+    split: str | None
+    """The dataset split the item belongs to (e.g. train, val, test)."""
+
+    metadata: Dict[str, Any] | None = None
+    """Dictionary holding additional metadata."""
+
+
+@dataclasses.dataclass
+class ITEM_DICT_ENTRY:
+    """Typing for holding queue items and number of save operations."""
+
+    items: List[QUEUE_ITEM]
+    """List of queue items."""
+
+    save_count: int
+    """Number of prior item batch saves to same file."""

eva/core/data/datasets/__init__.py

@@ -1,11 +1,11 @@
 """Datasets API."""
 
 from eva.core.data.datasets.base import Dataset
-from eva.core.data.datasets.dataset import TorchDataset
-from eva.core.data.datasets.embeddings import (
+from eva.core.data.datasets.classification import (
     EmbeddingsClassificationDataset,
     MultiEmbeddingsClassificationDataset,
 )
+from eva.core.data.datasets.dataset import TorchDataset
 
 __all__ = [
     "Dataset",

eva/core/data/datasets/classification/__init__.py

@@ -0,0 +1,8 @@
+"""Embedding cllassification datasets API."""
+
+from eva.core.data.datasets.classification.embeddings import EmbeddingsClassificationDataset
+from eva.core.data.datasets.classification.multi_embeddings import (
+    MultiEmbeddingsClassificationDataset,
+)
+
+__all__ = ["EmbeddingsClassificationDataset", "MultiEmbeddingsClassificationDataset"]

eva/core/data/datasets/classification/embeddings.py

@@ -0,0 +1,34 @@
+"""Embeddings classification dataset."""
+
+import os
+
+import torch
+from typing_extensions import override
+
+from eva.core.data.datasets import embeddings as embeddings_base
+
+
+class EmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[torch.Tensor]):
+    """Embeddings dataset class for classification tasks."""
+
+    @override
+    def _load_embeddings(self, index: int) -> torch.Tensor:
+        filename = self.filename(index)
+        embeddings_path = os.path.join(self._root, filename)
+        tensor = torch.load(embeddings_path, map_location="cpu")
+        if isinstance(tensor, list):
+            if len(tensor) > 1:
+                raise ValueError(
+                    f"Expected a single tensor in the .pt file, but found {len(tensor)}."
+                )
+            tensor = tensor[0]
+        return tensor.squeeze(0)
+
+    @override
+    def _load_target(self, index: int) -> torch.Tensor:
+        target = self._data.at[index, self._column_mapping["target"]]
+        return torch.tensor(target, dtype=torch.int64)
+
+    @override
+    def __len__(self) -> int:
+        return len(self._data)

eva/core/data/datasets/{embeddings/classification → classification}/multi_embeddings.py

@@ -7,10 +7,10 @@ import numpy as np
 import torch
 from typing_extensions import override
 
-from eva.core.data.datasets.embeddings import base
+from eva.core.data.datasets import embeddings as embeddings_base
 
 
-class MultiEmbeddingsClassificationDataset(base.EmbeddingsDataset):
+class MultiEmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[torch.Tensor]):
     """Dataset class for where a sample corresponds to multiple embeddings.
 
     Example use case: Slide level dataset where each slide has multiple patch embeddings.
@@ -21,7 +21,7 @@ class MultiEmbeddingsClassificationDataset(base.EmbeddingsDataset):
         root: str,
         manifest_file: str,
         split: Literal["train", "val", "test"],
-        column_mapping: Dict[str, str] = base.default_column_mapping,
+        column_mapping: Dict[str, str] = embeddings_base.default_column_mapping,
         embeddings_transforms: Callable | None = None,
         target_transforms: Callable | None = None,
     ):
@@ -32,9 +32,9 @@ class MultiEmbeddingsClassificationDataset(base.EmbeddingsDataset):
         The manifest must have a `column_mapping["multi_id"]` column that contains the
         unique identifier group of embeddings. For oncology datasets, this would be usually
         the slide id. Each row in the manifest file points to a .pt file that can contain
-        one or multiple embeddings. There can also be multiple rows for the same `multi_id`,
-        in which case the embeddings from the different .pt files corresponding to that same
-        `multi_id` will be stacked along the first dimension.
+        one or multiple embeddings (either as a list or stacked tensors). There can also be
+        multiple rows for the same `multi_id`, in which case the embeddings from the different
+        .pt files corresponding to that same `multi_id` will be stacked along the first dimension.
 
         Args:
             root: Root directory of the dataset.
@@ -73,10 +73,14 @@ class MultiEmbeddingsClassificationDataset(base.EmbeddingsDataset):
         embedding_paths = self._data.loc[
             self._data[self._column_mapping["multi_id"]] == multi_id, self._column_mapping["path"]
         ].to_list()
-        embedding_paths = [os.path.join(self._root, path) for path in embedding_paths]
 
         # Load embeddings and stack them accross the first dimension
-        embeddings = [torch.load(path, map_location="cpu") for path in embedding_paths]
+        embeddings = []
+        for path in embedding_paths:
+            embedding = torch.load(os.path.join(self._root, path), map_location="cpu")
+            if isinstance(embedding, list):
+                embedding = torch.stack(embedding, dim=0)
+            embeddings.append(embedding.unsqueeze(0) if embedding.ndim == 1 else embedding)
         embeddings = torch.cat(embeddings, dim=0)
 
         if not embeddings.ndim == 2:
@@ -103,4 +107,4 @@ class MultiEmbeddingsClassificationDataset(base.EmbeddingsDataset):
 
     @override
     def __len__(self) -> int:
-        return len(self._data)
+        return len(self._multi_ids)