kaiko-eva 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

eva/language/models/wrappers/litellm.py

@@ -0,0 +1,77 @@
+"""LLM wrapper for litellm models."""
+
+from typing import Any, Dict, List
+
+from litellm import batch_completion  # type: ignore
+from loguru import logger
+from typing_extensions import override
+
+from eva.core.models.wrappers import base
+
+
+class LiteLLMTextModel(base.BaseModel[List[str], List[str]]):
+    """Wrapper class for using litellm for chat-based text generation.
+
+    This wrapper uses litellm's `completion` function which accepts a list of
+    message dicts. The `forward` method converts a string prompt into a chat
+    message with a default "user" role, optionally prepends a system message,
+    and includes an API key if provided.
+    """
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        model_kwargs: Dict[str, Any] | None = None,
+    ) -> None:
+        """Initializes the litellm chat model wrapper.
+
+        Args:
+            model_name_or_path: The model identifier (or name) for litellm
+                (e.g.,"openai/gpt-4o" or "anthropic/claude-3-sonnet-20240229").
+            model_kwargs: Additional keyword arguments to pass during
+                generation (e.g., `temperature`, `max_tokens`).
+        """
+        super().__init__()
+        self._model_name_or_path = model_name_or_path
+        self._model_kwargs = model_kwargs or {}
+        self.load_model()
+
+    @override
+    def load_model(self) -> None:
+        """Prepares the litellm model.
+
+        Note:
+            litellm doesn't require an explicit loading step; models are called
+            directly during generation. This method exists for API consistency.
+        """
+        pass
+
+    @override
+    def model_forward(self, prompts: List[str]) -> List[str]:
+        """Generates text using litellm.
+
+        Args:
+            prompts: A list of prompts to be converted into a "user" message.
+
+        Returns:
+            A list of generated text responses. Failed generations will contain
+            error messages instead of generated text.
+        """
+        messages = [[{"role": "user", "content": prompt}] for prompt in prompts]
+
+        responses = batch_completion(
+            model=self._model_name_or_path,
+            messages=messages,
+            **self._model_kwargs,
+        )
+
+        results = []
+        for i, response in enumerate(responses):
+            if isinstance(response, Exception):
+                error_msg = f"Error generating text for prompt {i}: {response}"
+                logger.error(error_msg)
+                raise RuntimeError(error_msg)
+            else:
+                results.append(response["choices"][0]["message"]["content"])
+
+        return results

eva/language/models/wrappers/vllm.py

@@ -0,0 +1,149 @@
+"""LLM wrapper for vLLM models."""
+
+from typing import Any, Dict, List, Sequence
+
+from loguru import logger
+from typing_extensions import override
+
+try:
+    from vllm import LLM, SamplingParams  # type: ignore
+    from vllm.inputs import TokensPrompt  # type: ignore
+    from vllm.transformers_utils.tokenizer import AnyTokenizer  # type: ignore
+except ImportError as e:
+    raise ImportError(
+        "vLLM is required for VLLMTextModel but not installed. "
+        "vLLM must be installed manually as it requires CUDA and is not included in dependencies. "
+        "Install with: pip install vllm "
+        "Note: vLLM requires Linux with CUDA support for optimal performance. "
+        "For alternatives, consider using HuggingFaceTextModel or LiteLLMTextModel."
+    ) from e
+
+from eva.core.models.wrappers import base
+
+
+class VLLMTextModel(base.BaseModel):
+    """Wrapper class for using vLLM for text generation.
+
+    This wrapper loads a vLLM model, sets up the tokenizer and sampling
+    parameters, and uses a chat template to convert a plain string prompt
+    into the proper input format for vLLM generation. It then returns the
+    generated text response.
+    """
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        model_kwargs: Dict[str, Any] | None = None,
+        generation_kwargs: Dict[str, Any] | None = None,
+    ) -> None:
+        """Initializes the vLLM model wrapper.
+
+        Args:
+            model_name_or_path: The model identifier (e.g., a Hugging Face
+             repo ID or local path).
+            model_kwargs: Arguments required to initialize the vLLM model,
+                see [link](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/llm.py)
+                for more information.
+            generation_kwargs: Arguments required to generate the output,
+                need to align with the arguments of
+                [vllm.SamplingParams](https://github.com/vllm-project/vllm/blob/main/vllm/sampling_params.py).
+
+        """
+        super().__init__()
+        self._model_name_or_path = model_name_or_path
+        self._model_kwargs = model_kwargs or {}
+        self._generation_kwargs = generation_kwargs or {}
+
+        # Postpone heavy LLM initialisation to avoid pickling issues
+        self._llm_model: LLM | None = None
+        self._llm_tokenizer: AnyTokenizer | None = None
+
+    @override
+    def load_model(self) -> None:
+        """Create the vLLM engine on first use.
+
+        This lazy initialisation keeps the wrapper picklable by Ray / Lightning.
+        """
+        if self._llm_model is not None:
+            return
+        self._llm_model = LLM(model=self._model_name_or_path, **self._model_kwargs)
+        if self._llm_model is None:
+            raise RuntimeError("Model not initialized")
+        self._llm_tokenizer = self._llm_model.get_tokenizer()
+
+    def _apply_chat_template(self, prompts: Sequence[str]) -> list[TokensPrompt]:
+        """Apply chat template to the messages.
+
+        Args:
+            prompts: List of raw user strings.
+
+        Returns:
+            List of encoded messages.
+
+        Raises:
+            ValueError: If the tokenizer does not have a chat template.
+        """
+        self.load_model()
+        if self._llm_tokenizer is None:
+            raise RuntimeError("Tokenizer not initialized")
+
+        if not hasattr(self._llm_tokenizer, "chat_template"):
+            raise ValueError("Tokenizer does not have a chat template.")
+
+        chat_messages = [[{"role": "user", "content": p}] for p in prompts]
+        encoded_messages = self._llm_tokenizer.apply_chat_template(
+            chat_messages,  # type: ignore
+            tokenize=True,
+            add_generation_prompt=True,
+        )
+
+        # Check for double start token (BOS)
+        if (
+            hasattr(self._llm_tokenizer, "bos_token_id")
+            and self._llm_tokenizer.bos_token_id is not None
+            and isinstance(encoded_messages, list)
+            and len(encoded_messages) >= 2
+            and encoded_messages[0] == self._llm_tokenizer.bos_token_id
+            and encoded_messages[1] == self._llm_tokenizer.bos_token_id
+        ):
+
+            logger.warning("Found a double start token in the input_ids. Removing it.")
+            encoded_messages = encoded_messages[1:]
+
+        result = []
+        for encoded_message in encoded_messages:
+            if isinstance(encoded_message, (list, tuple)):
+                # Ensure all elements are integers
+                token_ids = [
+                    int(token) if isinstance(token, (int, str)) and str(token).isdigit() else 0
+                    for token in encoded_message
+                ]
+            else:
+                # Handle single token case
+                token_id = (
+                    int(encoded_message)
+                    if isinstance(encoded_message, (int, str)) and str(encoded_message).isdigit()
+                    else 0
+                )
+                token_ids = [token_id]
+
+            result.append(TokensPrompt(prompt_token_ids=token_ids))
+
+        return result
+
+    def generate(self, prompts: List[str]) -> List[str]:
+        """Generates text for the given prompt using the vLLM model.
+
+        Args:
+            prompts: A list of string prompts for generation.
+
+        Returns:
+            The generated text response.
+        """
+        self.load_model()
+        if self._llm_model is None:
+            raise RuntimeError("Model not initialized")
+
+        prompt_tokens = self._apply_chat_template(prompts)
+        outputs = self._llm_model.generate(prompt_tokens, SamplingParams(**self._generation_kwargs))
+        return [output.outputs[0].text for output in outputs]

eva/language/utils/__init__.py

@@ -0,0 +1,5 @@
+"""Language utilities and helper functions."""
+
+from eva.language.utils.str_to_int_tensor import CastStrToIntTensor
+
+__all__ = ["CastStrToIntTensor"]

eva/language/utils/str_to_int_tensor.py

@@ -0,0 +1,95 @@
+"""Transform utilities for post-processing predictions."""
+
+import re
+from typing import Any, Dict, List, Union
+
+import torch
+
+
+class CastStrToIntTensor:
+    """Casts string predictions to a torch.Tensor of ints using regex mapping.
+
+    This transform is useful when model outputs are text responses (e.g., 'yes', 'no', 'maybe')
+    that need to be converted into integer tensors for evaluation. It uses regex patterns
+    to map text responses to integer labels, making it flexible for various classification tasks.
+
+    Supports single values, lists of strings, or lists of integers.
+
+    Example:
+        >>> # Default mapping for yes/no/maybe classification
+        >>> transform = CastStrToIntTensor()
+        >>> transform(['yes', 'no', 'maybe'])
+        tensor([1, 0, 2])
+        >>> transform('yes')
+        tensor([1])
+
+        >>> # Custom mapping
+        >>> transform = CastStrToIntTensor({r'positive|good': 1, r'negative|bad': 0})
+        >>> transform(['positive', 'bad'])
+        tensor([1, 0])
+    """
+
+    def __init__(self, mapping: Dict[str, int] | None = None):
+        """Initialize the transform with a regex-to-integer mapping.
+
+        Args:
+            mapping: Dictionary mapping regex patterns to integers. If None, uses default
+                    yes/no/maybe mapping: {'no': 0, 'yes': 1, 'maybe': 2}
+        """
+        if mapping is None:
+            self.mapping = {r"\bno\b": 0, r"\byes\b": 1, r"\bmaybe\b": 2}
+        else:
+            self.mapping = mapping
+
+        self.compiled_patterns = [
+            (re.compile(pattern, re.IGNORECASE), value) for pattern, value in self.mapping.items()
+        ]
+
+    def __call__(self, values: Union[str, List[str], List[int]]) -> torch.Tensor:
+        """Convert string or list of strings/ints to a torch.Tensor of integers.
+
+        Args:
+            values: A string, or a list of strings/integers representing responses.
+
+        Returns:
+            A 1D torch.Tensor of integers.
+
+        Raises:
+            ValueError: If any value cannot be mapped to an integer.
+        """
+        return torch.tensor(
+            [self._cast_single(v) for v in (values if isinstance(values, list) else [values])],
+            dtype=torch.int,
+        )
+
+    def _cast_single(self, value: Any) -> int:
+        """Casts a single value to an integer using regex mapping.
+
+        Args:
+            value: A single value to convert (typically a string or int).
+
+        Returns:
+            The value as an integer.
+
+        Raises:
+            ValueError: If the value cannot be mapped to an integer.
+        """
+        if isinstance(value, int):
+            return value
+
+        if not isinstance(value, str):
+            value = str(value)
+
+        value = value.strip()
+
+        for pattern, mapped_value in self.compiled_patterns:
+            if pattern.search(value):
+                return mapped_value
+
+        try:
+            return int(value)
+        except (ValueError, TypeError) as e:
+            raise ValueError(
+                f"Cannot map value to int: {value!r}. "
+                f"Available patterns: {list(self.mapping.keys())}"
+            ) from e

eva/vision/data/dataloaders/__init__.py

@@ -1,5 +1,6 @@
 """Dataloader related utilities and functions."""
 
 from eva.vision.data.dataloaders import collate_fn
+from eva.vision.data.dataloaders.worker_init import seed_worker
 
-__all__ = ["collate_fn"]
+__all__ = ["collate_fn", "seed_worker"]

eva/vision/data/dataloaders/worker_init.py

@@ -0,0 +1,35 @@
+"""Dataloader worker init functions."""
+
+import random
+
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision.transforms.v2
+
+from eva.vision.data.transforms import base
+
+
+def seed_worker(worker_id: int) -> None:
+    """Sets the random seed for each dataloader worker process.
+
+    How to use?
+    `torch.utils.data.Dataloader(..., worker_init_fn=seed_worker)`
+
+    Args:
+        worker_id: The ID of the worker process.
+    """
+    worker_seed = (torch.initial_seed() + worker_id) % 2**32
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+    torch.manual_seed(worker_seed)
+
+    worker_info = torch.utils.data.get_worker_info()
+    if worker_info is not None and hasattr(worker_info, "dataset"):
+        dataset = torch.utils.data.get_worker_info().dataset  # type: ignore
+        if hasattr(dataset, "_transforms"):
+            transforms = dataset._transforms  # type: ignore
+            if isinstance(transforms, torchvision.transforms.v2.Compose):
+                for transform in transforms.transforms:
+                    if isinstance(transform, base.RandomMonaiTransform):
+                        transform.set_random_state(seed=worker_seed)

eva/vision/data/datasets/__init__.py

@@ -19,9 +19,9 @@ from eva.vision.data.datasets.segmentation import (
     BTCV,
     CoNSeP,
     EmbeddingsSegmentationDataset,
-    LiTS,
-    LiTSBalanced,
+    LiTS17,
     MoNuSAC,
+    MSDTask7Pancreas,
     TotalSegmentator2D,
 )
 from eva.vision.data.datasets.vision import VisionDataset
@@ -40,14 +40,14 @@ __all__ = [
     "PANDASmall",
     "Camelyon16",
     "PatchCamelyon",
+    "TotalSegmentator2D",
     "UniToPatho",
     "WsiClassificationDataset",
     "CoNSeP",
     "EmbeddingsSegmentationDataset",
-    "LiTS",
-    "LiTSBalanced",
+    "LiTS17",
+    "MSDTask7Pancreas",
     "MoNuSAC",
-    "TotalSegmentator2D",
     "VisionDataset",
     "MultiWsiDataset",
     "WsiDataset",

eva/vision/data/datasets/segmentation/__init__.py

@@ -4,9 +4,9 @@ from eva.vision.data.datasets.segmentation.bcss import BCSS
 from eva.vision.data.datasets.segmentation.btcv import BTCV
 from eva.vision.data.datasets.segmentation.consep import CoNSeP
 from eva.vision.data.datasets.segmentation.embeddings import EmbeddingsSegmentationDataset
-from eva.vision.data.datasets.segmentation.lits import LiTS
-from eva.vision.data.datasets.segmentation.lits_balanced import LiTSBalanced
+from eva.vision.data.datasets.segmentation.lits17 import LiTS17
 from eva.vision.data.datasets.segmentation.monusac import MoNuSAC
+from eva.vision.data.datasets.segmentation.msd_task7_pancreas import MSDTask7Pancreas
 from eva.vision.data.datasets.segmentation.total_segmentator_2d import TotalSegmentator2D
 
 __all__ = [
@@ -14,8 +14,8 @@ __all__ = [
     "BTCV",
     "CoNSeP",
     "EmbeddingsSegmentationDataset",
-    "LiTS",
-    "LiTSBalanced",
+    "LiTS17",
+    "MSDTask7Pancreas",
     "MoNuSAC",
     "TotalSegmentator2D",
 ]

eva/vision/data/datasets/segmentation/btcv.py

@@ -10,6 +10,7 @@ from torchvision import tv_tensors
 from torchvision.datasets import utils as data_utils
 from typing_extensions import override
 
+from eva.core.utils import requirements
 from eva.vision.data import tv_tensors as eva_tv_tensors
 from eva.vision.data.datasets import _utils as _data_utils
 from eva.vision.data.datasets.segmentation import _utils
@@ -105,6 +106,8 @@ class BTCV(VisionDataset[eva_tv_tensors.Volume, tv_tensors.Mask]):
 
     @override
     def validate(self) -> None:
+        requirements.check_dependencies(requirements={"torch": "2.5.1", "torchvision": "0.20.1"})
+
         def _valid_sample(index: int) -> bool:
             """Indicates if the sample files exist and are reachable."""
             volume_file, segmentation_file = self._samples[self._indices[index]]

eva/vision/data/datasets/segmentation/consep.py

@@ -108,6 +108,11 @@ class CoNSeP(wsi.MultiWsiDataset, vision.VisionDataset[tv_tensors.Image, tv_tens
             n_classes=5,
             first_and_last_labels=((self.classes[0], self.classes[-1])),
         )
+        n_expected = self._expected_dataset_lengths[None]
+        if len(self._file_paths) != n_expected:
+            raise ValueError(
+                f"Expected {n_expected} images, found {len(self._file_paths)} in {self._root}."
+            )
 
     @override
     def __getitem__(self, index: int) -> Tuple[tv_tensors.Image, tv_tensors.Mask, Dict[str, Any]]:
@@ -135,10 +140,6 @@ class CoNSeP(wsi.MultiWsiDataset, vision.VisionDataset[tv_tensors.Image, tv_tens
     def _load_file_paths(self, split: Literal["train", "val"] | None = None) -> List[str]:
         """Loads the file paths of the corresponding dataset split."""
         paths = list(glob.glob(os.path.join(self._root, "**/Images/*.png"), recursive=True))
-        n_expected = self._expected_dataset_lengths[None]
-        if len(paths) != n_expected:
-            raise ValueError(f"Expected {n_expected} images, found {len(paths)} in {self._root}.")
-
         if split is not None:
             split_to_folder = {"train": "Train", "val": "Test"}
             paths = filter(lambda p: split_to_folder[split] == p.split("/")[-3], paths)