kaiko-eva 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

eva/language/utils/text/messages.py

@@ -0,0 +1,113 @@
+"""Message formatting utilities for language models."""
+
+import functools
+import json
+from typing import Any, Dict, List
+
+from eva.language.data.messages import (
+    AssistantMessage,
+    MessageSeries,
+    Role,
+    SystemMessage,
+    UserMessage,
+)
+
+
+def format_chat_message(message: MessageSeries) -> List[Dict[str, Any]]:
+    """Formats a message series into a format following OpenAI's API specification."""
+    return [{"role": item.role, "content": item.content} for item in message]
+
+
+def combine_system_messages(message: MessageSeries, join_char: str = "\n") -> MessageSeries:
+    """Combine system messages into a single message.
+
+    This is useful when the MessageSeries contains multiple system messages such
+    as `ModelSystemMessage` and `TaskSystemMessage`. But given that most models / apis
+    expect a single system message, this function can be used to combines them into one.
+
+    Args:
+        message: The message series containing one or multiple messages.
+        join_char: The character to use to join the system messages. Default is newline.
+
+    Returns:
+        A new message series with system messages combined into one and the
+        remaining messages unchanged.
+    """
+    system_messages = list(filter(lambda item: item.role == Role.SYSTEM, message))
+    if len(system_messages) == 0:
+        return message
+
+    non_system_messages = list(filter(lambda item: item.role != Role.SYSTEM, message))
+    return [
+        SystemMessage(content=merge_message_contents(system_messages, join_char=join_char))
+    ] + non_system_messages
+
+
+def merge_message_contents(message: MessageSeries, join_char: str = "\n") -> str:
+    """Merges the all contents within a message series into a string.
+
+    Args:
+        message: The message series to combine.
+        join_char: The character to use to join the message contents. Default is newline.
+
+    Returns:
+        A string containing the combined message contents.
+    """
+    return join_char.join(item.content for item in message)
+
+
+def insert_system_message(
+    message: MessageSeries, system_message: SystemMessage | None
+) -> MessageSeries:
+    """Insert a system message at the beginning of the message series."""
+    if system_message is None:
+        return message
+    return [system_message] + message
+
+
+def batch_insert_system_message(
+    messages: List[MessageSeries], system_message: SystemMessage | None
+) -> List[MessageSeries]:
+    """Insert a system message at the beginning of each message series in a batch."""
+    return list(
+        map(functools.partial(insert_system_message, system_message=system_message), messages)
+    )
+
+
+def serialize(messages: MessageSeries) -> str:
+    """Serialize a MessageSeries object into a JSON string.
+
+    Args:
+        messages: A list of message objects (MessagesSeries).
+
+    Returns:
+        A JSON string representing the message series, with the following format:
+            [{"role": "user", "content": "Hello"}, ...]
+    """
+    serialized_messages = format_chat_message(messages)
+    return json.dumps(serialized_messages)
+
+
+def deserialize(messages: str) -> MessageSeries:
+    """Convert a json string to a MessageSeries object.
+
+    Format: [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}]
+    """
+    message_dicts = json.loads(messages)
+
+    message_series = []
+    for message_dict in message_dicts:
+        if "role" not in message_dict or "content" not in message_dict:
+            raise ValueError("`role` or `content` keys are missing.")
+
+        match message_dict["role"]:
+            case Role.USER:
+                message_series.append(UserMessage(**message_dict))
+            case Role.ASSISTANT:
+                message_series.append(AssistantMessage(**message_dict))
+            case Role.SYSTEM:
+                message_series.append(SystemMessage(**message_dict))
+            case _:
+                raise ValueError(f"Unknown role: {message_dict['role']}")
+
+    return message_series

eva/multimodal/__init__.py

@@ -0,0 +1,6 @@
+"""Multimodal API."""
+
+from eva.multimodal import models
+from eva.multimodal.data import datasets
+
+__all__ = ["models", "datasets"]

eva/multimodal/callbacks/__init__.py

@@ -0,0 +1,5 @@
+"""Multimodal callbacks API."""
+
+from eva.multimodal.callbacks.writers import TextPredictionWriter
+
+__all__ = ["TextPredictionWriter"]

eva/multimodal/callbacks/writers/__init__.py

@@ -0,0 +1,5 @@
+"""Multimodal writers callbacks API."""
+
+from eva.multimodal.callbacks.writers.prediction import TextPredictionWriter
+
+__all__ = ["TextPredictionWriter"]

eva/multimodal/callbacks/writers/prediction.py

@@ -0,0 +1,39 @@
+"""Text prediction writer callbacks."""
+
+from typing import Dict, List, Literal, Tuple
+
+from torch import nn
+from typing_extensions import override
+
+from eva.language.callbacks import writers
+from eva.multimodal.models.typings import TextImageBatch
+
+
+class TextPredictionWriter(writers.TextPredictionWriter):
+    """Callback for writing generated text predictions to disk."""
+
+    def __init__(
+        self,
+        output_dir: str,
+        model: nn.Module,
+        dataloader_idx_map: Dict[int, str] | None = None,
+        metadata_keys: List[str] | None = None,
+        include_input: bool = True,
+        overwrite: bool = False,
+        save_format: Literal["jsonl", "parquet", "csv"] = "jsonl",
+    ) -> None:
+        """See docstring of base class."""
+        super().__init__(
+            output_dir=output_dir,
+            model=model,
+            dataloader_idx_map=dataloader_idx_map,
+            metadata_keys=metadata_keys,
+            include_input=include_input,
+            overwrite=overwrite,
+            save_format=save_format,
+        )
+
+    @override
+    def _unpack_batch(self, batch: TextImageBatch) -> Tuple[list, list | None, dict | None]:  # type: ignore
+        text_batch, _, target_batch, metadata_batch = TextImageBatch(*batch)
+        return text_batch, target_batch, metadata_batch

eva/multimodal/data/__init__.py

@@ -0,0 +1,5 @@
+"""Data components for multimodal learning."""
+
+from eva.multimodal.data import datasets
+
+__all__ = ["datasets"]

eva/multimodal/data/dataloaders/__init__.py

@@ -0,0 +1,5 @@
+"""Multimodal dataloaders API."""
+
+from eva.multimodal.data.dataloaders.collate_fn import text_image_collate
+
+__all__ = ["text_image_collate"]

eva/multimodal/data/dataloaders/collate_fn/__init__.py

@@ -0,0 +1,5 @@
+"""Multimodal collate functions API."""
+
+from eva.multimodal.data.dataloaders.collate_fn.text_image import text_image_collate
+
+__all__ = ["text_image_collate"]

eva/multimodal/data/dataloaders/collate_fn/text_image.py

@@ -0,0 +1,28 @@
+"""Collate functions for text-image data."""
+
+from typing import List
+
+from torch.utils.data._utils.collate import default_collate
+
+from eva.multimodal.data.datasets.typings import TextImageSample
+from eva.multimodal.models.typings import TextImageBatch
+
+
+def text_image_collate(batch: List[TextImageSample]) -> TextImageBatch:
+    """Collate function for text-image batches."""
+    texts, images, targets, metadata = zip(*batch, strict=False)
+
+    first_sample = batch[0]
+    metadata = None
+    if first_sample.metadata is not None:
+        metadata = {
+            k: [sample.metadata[k] for sample in batch if sample.metadata]
+            for k in first_sample.metadata.keys()
+        }
+
+    return TextImageBatch(
+        text=list(texts),
+        image=list(images),
+        target=default_collate(targets) if targets[0] is not None else None,
+        metadata=metadata,
+    )

eva/multimodal/data/datasets/__init__.py

@@ -0,0 +1,6 @@
+"""Multimodal datasets API."""
+
+from eva.multimodal.data.datasets.multiple_choice.patch_camelyon import PatchCamelyon
+from eva.multimodal.data.datasets.text_image import TextImageDataset
+
+__all__ = ["TextImageDataset", "PatchCamelyon"]

eva/multimodal/data/datasets/base.py

@@ -0,0 +1,13 @@
+"""Multimodal Dataset base class."""
+
+import abc
+from typing import Generic, TypeVar
+
+from eva.core.data.datasets import base
+
+DataSample = TypeVar("DataSample")
+"""The data sample type."""
+
+
+class MultimodalDataset(base.MapDataset, abc.ABC, Generic[DataSample]):
+    """Base dataset class for multimodal tasks."""

eva/multimodal/data/datasets/multiple_choice/__init__.py

@@ -0,0 +1,5 @@
+"""Multiple choice datasets."""
+
+from eva.multimodal.data.datasets.multiple_choice.patch_camelyon import PatchCamelyon
+
+__all__ = ["PatchCamelyon"]

eva/multimodal/data/datasets/multiple_choice/patch_camelyon.py

@@ -0,0 +1,80 @@
+"""PatchCamelyon dataset with text prompts for multimodal classification."""
+
+from typing import Any, Dict, Literal
+
+from torchvision import tv_tensors
+from typing_extensions import override
+
+from eva.language.data.messages import MessageSeries, UserMessage
+from eva.multimodal.data.datasets.schemas import TransformsSchema
+from eva.multimodal.data.datasets.text_image import TextImageDataset
+from eva.vision.data import datasets as vision_datasets
+
+
+class PatchCamelyon(TextImageDataset[int], vision_datasets.PatchCamelyon):
+    """PatchCamelyon image classification using a multiple choice text prompt."""
+
+    _default_prompt = (
+        "You are a pathology expert helping pathologists to analyze images of tissue samples.\n"
+        "Question: Does this image show metastatic breast tissue?\n"
+        "Options: A: no, B: yes\n"
+        "Only answer with a single letter without further explanation. "
+        "Please always provide an answer, even if you are not sure.\n"
+        "Answer: "
+    )
+
+    def __init__(
+        self,
+        root: str,
+        split: Literal["train", "val", "test"],
+        download: bool = False,
+        transforms: TransformsSchema | None = None,
+        prompt: str | None = None,
+        max_samples: int | None = None,
+    ) -> None:
+        """Initializes the dataset.
+
+        Args:
+            root: The path to the dataset root. This path should contain
+                the uncompressed h5 files and the metadata.
+            split: The dataset split for training, validation, or testing.
+            download: Whether to download the data for the specified split.
+                Note that the download will be executed only by additionally
+                calling the :meth:`prepare_data` method.
+            transforms: A function/transform which returns a transformed
+                version of the raw data samples.
+            prompt: The text prompt to use for classification (multple choice).
+            max_samples: Maximum number of samples to use. If None, use all samples.
+        """
+        super().__init__(root=root, split=split, download=download, transforms=transforms)
+
+        self.max_samples = max_samples
+        self.prompt = prompt or self._default_prompt
+
+        if self.max_samples is not None:
+            self._expected_length = {split: max_samples}
+
+    @property
+    @override
+    def class_to_idx(self) -> Dict[str, int]:
+        return {"A": 0, "B": 1}
+
+    @override
+    def __len__(self) -> int:
+        return self.max_samples or self._fetch_dataset_length()
+
+    @override
+    def load_text(self, index: int) -> MessageSeries:
+        return [UserMessage(content=self.prompt)]
+
+    @override
+    def load_image(self, index: int) -> tv_tensors.Image:
+        return vision_datasets.PatchCamelyon.load_data(self, index)
+
+    @override
+    def load_target(self, index: int) -> int:
+        return int(vision_datasets.PatchCamelyon.load_target(self, index).item())
+
+    @override
+    def load_metadata(self, index: int) -> Dict[str, Any] | None:
+        return vision_datasets.PatchCamelyon.load_metadata(self, index)

eva/multimodal/data/datasets/schemas.py

@@ -0,0 +1,14 @@
+"""Schema definitions for dataset classes."""
+
+import dataclasses
+from typing import Callable
+
+from eva.language.data.datasets import schemas as language_schemas
+
+
+@dataclasses.dataclass(frozen=True)
+class TransformsSchema(language_schemas.TransformsSchema):
+    """Schema for dataset transforms."""
+
+    image: Callable | None = None
+    """Image transformation"""

eva/multimodal/data/datasets/text_image.py

@@ -0,0 +1,77 @@
+"""Base classes for text-image datasets."""
+
+import abc
+from typing import Generic
+
+from torchvision import tv_tensors
+from typing_extensions import override
+
+from eva.language.data.datasets.text import TextDataset
+from eva.multimodal.data.datasets.base import MultimodalDataset
+from eva.multimodal.data.datasets.schemas import TransformsSchema
+from eva.multimodal.data.datasets.typings import TargetType, TextImageSample
+
+
+class TextImageDataset(
+    MultimodalDataset[TextImageSample[TargetType]], TextDataset, abc.ABC, Generic[TargetType]
+):
+    """Base dataset class for text-image tasks."""
+
+    def __init__(self, *args, transforms: TransformsSchema | None = None, **kwargs) -> None:
+        """Initializes the dataset.
+
+        Args:
+            *args: Positional arguments for the base class.
+            transforms: The transforms to apply to the text, image and target when
+                loading the samples.
+            **kwargs: Keyword arguments for the base class.
+        """
+        super().__init__(*args, **kwargs)
+
+        self.transforms = transforms
+
+    @abc.abstractmethod
+    def load_image(self, index: int) -> tv_tensors.Image:
+        """Returns the image content.
+
+        Args:
+            index: The index of the data sample.
+
+        Returns:
+            The image content.
+        """
+        raise NotImplementedError
+
+    @override
+    def __getitem__(self, index: int) -> TextImageSample[TargetType]:
+        item = TextImageSample(
+            text=self.load_text(index),
+            image=self.load_image(index),
+            target=self.load_target(index),
+            metadata=self.load_metadata(index) or {},
+        )
+        return self._apply_transforms(item)
+
+    @override
+    def _apply_transforms(self, sample: TextImageSample[TargetType]) -> TextImageSample[TargetType]:
+        """Applies the dataset transforms to the text, image and target.
+
+        Args:
+            sample: The sample containing text, image, target and metadata.
+
+        Returns:
+            The transformed sample.
+        """
+        if self.transforms:
+            text = self.transforms.text(sample.text) if self.transforms.text else sample.text
+            image = self.transforms.image(sample.image) if self.transforms.image else sample.image
+            target = (
+                self.transforms.target(sample.target) if self.transforms.target else sample.target
+            )
+            return TextImageSample(
+                text=text,
+                image=image,
+                target=target,
+                metadata=sample.metadata,
+            )
+        return sample

eva/multimodal/data/datasets/typings.py

@@ -0,0 +1,27 @@
+"""Typings for multimodal datasets."""
+
+from typing import Any, Generic, TypeVar
+
+from torchvision import tv_tensors
+from typing_extensions import NamedTuple
+
+from eva.language.data.messages import MessageSeries
+
+TargetType = TypeVar("TargetType")
+"""The target data type."""
+
+
+class TextImageSample(NamedTuple, Generic[TargetType]):
+    """Text and image sample with target and metadata."""
+
+    text: MessageSeries
+    """One or multiple conversation messages."""
+
+    image: tv_tensors.Image
+    """Image tensor."""
+
+    target: TargetType | None
+    """Target data."""
+
+    metadata: dict[str, Any] | None
+    """Additional metadata."""

eva/multimodal/models/__init__.py

@@ -0,0 +1,8 @@
+"""Multimodal models API."""
+
+from eva.multimodal.models import networks, wrappers
+
+__all__ = [
+    "networks",
+    "wrappers",
+]

eva/multimodal/models/modules/__init__.py

@@ -0,0 +1,5 @@
+"""Multimodal Networks API."""
+
+from eva.multimodal.models.modules.vision_language import VisionLanguageModule
+
+__all__ = ["VisionLanguageModule"]

eva/multimodal/models/modules/vision_language.py

@@ -0,0 +1,55 @@
+"""Model module for vision-language models."""
+
+from typing import Any, List
+
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from torch import nn
+from typing_extensions import override
+
+from eva.core.metrics import structs as metrics_lib
+from eva.core.models.modules import module
+from eva.core.models.modules.utils import batch_postprocess
+from eva.multimodal.models.typings import TextImageBatch
+
+
+class VisionLanguageModule(module.ModelModule):
+    """Model module for vision-language tasks."""
+
+    def __init__(
+        self,
+        model: nn.Module,
+        metrics: metrics_lib.MetricsSchema | None = None,
+        postprocess: batch_postprocess.BatchPostProcess | None = None,
+    ) -> None:
+        """Initializes the text inference module.
+
+        Args:
+            model: Model instance to use for forward pass.
+            metrics: Metrics schema for evaluation.
+            postprocess: A helper function to post-process model outputs before evaluation.
+        """
+        super().__init__(metrics=metrics, postprocess=postprocess)
+
+        self.model = model
+
+    @override
+    def forward(self, batch: TextImageBatch, *args: Any, **kwargs: Any) -> List[str]:
+        return self.model(batch)
+
+    @override
+    def validation_step(self, batch: TextImageBatch, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    @override
+    def test_step(self, batch: TextImageBatch, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    def _batch_step(self, batch: TextImageBatch) -> STEP_OUTPUT:
+        text, _, targets, metadata = TextImageBatch(*batch)
+        predictions = self.forward(batch)
+        return {
+            "inputs": text,
+            "predictions": predictions,
+            "targets": targets,
+            "metadata": metadata,
+        }

eva/multimodal/models/networks/__init__.py

@@ -0,0 +1,14 @@
+"""Multimodal networks API."""
+
+from eva.multimodal.models.networks.alibaba import Qwen25VL7BInstruct
+from eva.multimodal.models.networks.api import Claude35Sonnet20240620, Claude37Sonnet20250219
+from eva.multimodal.models.networks.others import PathoR13b
+from eva.multimodal.models.networks.registry import model_registry
+
+__all__ = [
+    "Claude35Sonnet20240620",
+    "Claude37Sonnet20250219",
+    "PathoR13b",
+    "Qwen25VL7BInstruct",
+    "model_registry",
+]

eva/multimodal/models/networks/alibaba.py

@@ -0,0 +1,39 @@
+"""Models from Alibaba."""
+
+import torch
+
+from eva.multimodal.models import wrappers
+from eva.multimodal.models.networks.registry import model_registry
+
+
+@model_registry.register("alibaba/qwen2-5-vl-7b-instruct")
+class Qwen25VL7BInstruct(wrappers.HuggingFaceModel):
+    """Qwen2.5-VL 7B Instruct model."""
+
+    def __init__(
+        self,
+        system_prompt: str | None = None,
+        cache_dir: str | None = None,
+        attn_implementation: str = "flash_attention_2",
+    ):
+        """Initialize the model."""
+        super().__init__(
+            model_name_or_path="Qwen/Qwen2.5-VL-7B-Instruct",
+            model_class="Qwen2_5_VLForConditionalGeneration",
+            model_kwargs={
+                "torch_dtype": torch.bfloat16,
+                "trust_remote_code": True,
+                "cache_dir": cache_dir,
+                "attn_implementation": attn_implementation,
+            },
+            generation_kwargs={
+                "max_new_tokens": 512,
+                "do_sample": False,
+            },
+            processor_kwargs={
+                "padding": True,
+                "padding_side": "left",
+                "max_pixels": 451584,  # 672*672
+            },
+            system_prompt=system_prompt,
+        )

eva/multimodal/models/networks/api/__init__.py

@@ -0,0 +1,11 @@
+"""Multimodal API networks."""
+
+from eva.multimodal.models.networks.api.anthropic import (
+    Claude35Sonnet20240620,
+    Claude37Sonnet20250219,
+)
+
+__all__ = [
+    "Claude35Sonnet20240620",
+    "Claude37Sonnet20250219",
+]

eva/multimodal/models/networks/api/anthropic.py

@@ -0,0 +1,34 @@
+"""Models from Anthropic."""
+
+import os
+
+from eva.multimodal.models import wrappers
+from eva.multimodal.models.networks.registry import model_registry
+
+
+class _Claude(wrappers.LiteLLMModel):
+    """Base class for Claude models."""
+
+    def __init__(self, model_name: str, system_prompt: str | None = None):
+        if not os.getenv("ANTHROPIC_API_KEY"):
+            raise ValueError("ANTHROPIC_API_KEY env variable must be set.")
+
+        super().__init__(model_name=model_name, system_prompt=system_prompt)
+
+
+@model_registry.register("anthropic/claude-3-5-sonnet-20240620")
+class Claude35Sonnet20240620(_Claude):
+    """Claude 3.5 Sonnet (June 2024) model."""
+
+    def __init__(self, system_prompt: str | None = None):
+        """Initialize the model."""
+        super().__init__(model_name="claude-3-5-sonnet-20240620", system_prompt=system_prompt)
+
+
+@model_registry.register("anthropic/claude-3-7-sonnet-20250219")
+class Claude37Sonnet20250219(_Claude):
+    """Claude 3.7 Sonnet (February 2025) model."""
+
+    def __init__(self, system_prompt: str | None = None):
+        """Initialize the model."""
+        super().__init__(model_name="claude-3-7-sonnet-20250219", system_prompt=system_prompt)