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

eva/language/data/datasets/classification/pubmedqa.py

@@ -10,11 +10,20 @@ from loguru import logger
 from typing_extensions import override
 
 from eva.language.data.datasets.classification import base
+from eva.language.data.messages import MessageSeries, UserMessage
 
 
 class PubMedQA(base.TextClassification):
     """Dataset class for PubMedQA question answering task."""
 
+    _expected_dataset_lengths: Dict[str | None, int] = {
+        "train": 450,
+        "val": 50,
+        "test": 500,
+        None: 500,
+    }
+    """Expected dataset lengths for the splits and complete dataset."""
+
     _license: str = "MIT License (https://github.com/pubmedqa/pubmedqa/blob/master/LICENSE)"
     """Dataset license."""
 
@@ -52,7 +61,14 @@ class PubMedQA(base.TextClassification):
         """
         dataset_name = "bigbio/pubmed_qa"
         config_name = "pubmed_qa_labeled_fold0_source"
-        split = (self._split or "train+test+validation") if self._split != "val" else "validation"
+
+        match self._split:
+            case "val":
+                split = "validation"
+            case None:
+                split = "train+test+validation"
+            case _:
+                split = self._split
 
         if self._download:
             logger.info("Downloading dataset from HuggingFace Hub")
@@ -88,7 +104,7 @@ class PubMedQA(base.TextClassification):
         dataset_path = None
 
         if self._root:
-            dataset_path = self._root
+            dataset_path = os.path.join(self._root, self._split) if self._split else self._root
             os.makedirs(self._root, exist_ok=True)
 
         try:
@@ -103,6 +119,15 @@ class PubMedQA(base.TextClassification):
         except Exception as e:
             raise RuntimeError(f"Failed to prepare dataset: {e}") from e
 
+    @override
+    def validate(self) -> None:
+        if len(self) != self._expected_dataset_lengths[self._split]:
+            raise ValueError(
+                f"Dataset length mismatch for split '{self._split}': "
+                f"expected {self._expected_dataset_lengths[self._split]}, "
+                f"but got {len(self)}"
+            )
+
     @property
     @override
     def classes(self) -> List[str]:
@@ -114,11 +139,18 @@ class PubMedQA(base.TextClassification):
         return {"no": 0, "yes": 1, "maybe": 2}
 
     @override
-    def load_text(self, index: int) -> str:
+    def load_text(self, index: int) -> MessageSeries:
         if index < 0 or index >= len(self.dataset):
             raise IndexError(f"Index {index} out of range for dataset of size {len(self.dataset)}")
         sample = dict(self.dataset[index])
-        return f"Question: {sample['QUESTION']}\nContext: " + " ".join(sample["CONTEXTS"])
+        return [
+            UserMessage(
+                content=f"Question: {sample['QUESTION']}\nContext: "
+                + " ".join(sample["CONTEXTS"])
+                + "\nInstruction: Carefully read the question and the provided context. "
+                + "Answer with one word: 'yes', 'no', or 'maybe'. Answer: "
+            )
+        ]
 
     @override
     def load_target(self, index: int) -> torch.Tensor:

eva/language/data/datasets/prediction.py

@@ -0,0 +1,151 @@
+"""Dataset class for loading pre-generated text predictions."""
+
+import abc
+from pathlib import Path
+from typing import Any, Dict, Generic, Literal
+
+import pandas as pd
+from typing_extensions import override
+
+from eva.language.data.datasets.base import LanguageDataset
+from eva.language.data.datasets.schemas import TransformsSchema
+from eva.language.data.datasets.typings import PredictionSample, TargetType
+from eva.language.data.messages import MessageSeries, UserMessage
+from eva.language.utils.text import messages as message_utils
+
+
+class TextPredictionDataset(
+    LanguageDataset[PredictionSample[TargetType]], abc.ABC, Generic[TargetType]
+):
+    """Dataset class for loading pre-generated text predictions."""
+
+    def __init__(
+        self,
+        path: str,
+        prediction_column: str = "prediction",
+        target_column: str = "target",
+        text_column: str | None = None,
+        metadata_columns: list[str] | None = None,
+        split: Literal["train", "val", "test"] | None = None,
+        transforms: TransformsSchema | None = None,
+    ):
+        """Initialize the dataset.
+
+        Args:
+            path: The path to the manifest file holding the predictions & targets.
+            prediction_column: The name of the prediction column.
+            target_column: The name of the label column.
+            text_column: The name of the column with the text inputs that were used
+                to generate the predictions. If the text column contains chat message
+                json format ([{"role": ..., "content": ...}]), it will be deserialized into
+                a list of Message objects. Otherwise, the content is interpreted as a
+                single user message.
+            metadata_columns: List of column names to include in metadata.
+            split: The dataset split to use (train, val, test). If not specified,
+                the entire dataset will be used.
+            transforms: The transforms to apply to the text and target when
+                loading the samples.
+        """
+        super().__init__()
+
+        self.path = path
+        self.prediction_column = prediction_column
+        self.target_column = target_column
+        self.text_column = text_column
+        self.metadata_columns = metadata_columns
+        self.split = split
+        self.transforms = transforms
+
+        self._data: pd.DataFrame
+
+    @override
+    def __len__(self) -> int:
+        return len(self._data)
+
+    @override
+    def __getitem__(self, index: int) -> PredictionSample[TargetType]:
+        item = PredictionSample(
+            prediction=self.load_prediction(index),
+            target=self.load_target(index),
+            text=self.load_text(index),
+            metadata=self.load_metadata(index) or {},
+        )
+        return self._apply_transforms(item)
+
+    @override
+    def configure(self) -> None:
+        extension = Path(self.path).suffix
+
+        match extension:
+            case ".jsonl":
+                self._data = pd.read_json(self.path, lines=True)
+            case ".csv":
+                self._data = pd.read_csv(self.path)
+            case ".parquet":
+                self._data = pd.read_parquet(self.path)
+            case _:
+                raise ValueError(f"Unsupported file extension: {extension}")
+
+        if self.split is not None:
+            self._data = self._data[self._data["split"] == self.split].reset_index(drop=True)  # type: ignore
+
+    @override
+    def validate(self) -> None:
+        if self.prediction_column not in self._data.columns:
+            raise ValueError(f"Label column '{self.prediction_column}' not found.")
+        if self.target_column not in self._data.columns:
+            raise ValueError(f"Label column '{self.target_column}' not found.")
+        if self.metadata_columns:
+            missing_columns = set(self.metadata_columns) - set(self._data.columns)
+            if missing_columns:
+                raise ValueError(f"Metadata columns {missing_columns} not found.")
+
+    def load_prediction(self, index: int) -> TargetType:
+        """Returns the prediction for the given index."""
+        return self._data.iloc[index][self.prediction_column]
+
+    def load_target(self, index: int) -> TargetType:
+        """Returns the target for the given index."""
+        return self._data.iloc[index][self.target_column]
+
+    def load_text(self, index: int) -> MessageSeries | None:
+        """Returns the text for the given index."""
+        if self.text_column is None:
+            return None
+
+        text = self._data.iloc[index][self.text_column]
+
+        try:
+            return message_utils.deserialize(self._data.iloc[index][self.text_column])
+        except Exception:
+            return [UserMessage(content=text)]
+
+    def load_metadata(self, index: int) -> Dict[str, Any] | None:
+        """Returns the metadata for the given index."""
+        if self.metadata_columns is None:
+            return None
+
+        row = self._data.iloc[index]
+        return {col: row[col] for col in self.metadata_columns}
+
+    def _apply_transforms(
+        self, sample: PredictionSample[TargetType]
+    ) -> PredictionSample[TargetType]:
+        """Applies the dataset transforms to the prediction and target."""
+        if self.transforms:
+            text = self.transforms.text(sample.text) if self.transforms.text else sample.text
+            prediction = (
+                self.transforms.prediction(sample.prediction)
+                if self.transforms.prediction
+                else sample.prediction
+            )
+            target = (
+                self.transforms.target(sample.target) if self.transforms.target else sample.target
+            )
+            return PredictionSample(
+                prediction=prediction,
+                target=target,
+                text=text,
+                metadata=sample.metadata,
+            )
+        return sample

eva/language/data/datasets/schemas.py

@@ -0,0 +1,18 @@
+"""Schema definitions for dataset classes."""
+
+import dataclasses
+from typing import Callable
+
+
+@dataclasses.dataclass(frozen=True)
+class TransformsSchema:
+    """Schema for dataset transforms."""
+
+    text: Callable | None = None
+    """Text transformation"""
+
+    target: Callable | None = None
+    """Target transformation"""
+
+    prediction: Callable | None = None
+    """Prediction transformation"""

eva/language/data/datasets/text.py

@@ -0,0 +1,92 @@
+"""Base classes for text-image datasets."""
+
+import abc
+from typing import Any, Dict, Generic
+
+from typing_extensions import override
+
+from eva.language.data.datasets.base import LanguageDataset
+from eva.language.data.datasets.schemas import TransformsSchema
+from eva.language.data.datasets.typings import TargetType, TextSample
+from eva.language.data.messages import MessageSeries
+
+
+class TextDataset(LanguageDataset[TextSample[TargetType]], abc.ABC, Generic[TargetType]):
+    """Base dataset class for text-based 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 and target when
+                loading the samples.
+            **kwargs: Keyword arguments for the base class.
+        """
+        super().__init__(*args, **kwargs)
+
+        self.transforms = transforms
+
+    def load_metadata(self, index: int) -> Dict[str, Any] | None:
+        """Returns the dataset metadata.
+
+        Args:
+            index: The index of the data sample.
+
+        Returns:
+            The sample metadata.
+        """
+
+    @abc.abstractmethod
+    def load_text(self, index: int) -> MessageSeries:
+        """Returns the text content.
+
+        Args:
+            index: The index of the data sample.
+
+        Returns:
+            The text content.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def load_target(self, index: int) -> TargetType:
+        """Returns the target label.
+
+        Args:
+            index: The index of the data sample.
+
+        Returns:
+            The target label.
+        """
+        raise NotImplementedError
+
+    @override
+    def __getitem__(self, index: int) -> TextSample[TargetType]:
+        item = TextSample(
+            text=self.load_text(index),
+            target=self.load_target(index),
+            metadata=self.load_metadata(index) or {},
+        )
+        return self._apply_transforms(item)
+
+    def _apply_transforms(self, sample: TextSample[TargetType]) -> TextSample[TargetType]:
+        """Applies the dataset transforms to the text and target.
+
+        Args:
+            sample: The text sample..
+
+        Returns:
+            The transformed sample.
+        """
+        if self.transforms:
+            text = self.transforms.text(sample.text) if self.transforms.text else sample.text
+            target = (
+                self.transforms.target(sample.target) if self.transforms.target else sample.target
+            )
+            return TextSample(
+                text=text,
+                target=target,
+                metadata=sample.metadata,
+            )
+        return sample

eva/language/data/datasets/typings.py

@@ -0,0 +1,39 @@
+"""Typings for multimodal datasets."""
+
+from typing import Any, Generic, TypeVar
+
+from typing_extensions import NamedTuple
+
+from eva.language.data.messages import MessageSeries
+
+TargetType = TypeVar("TargetType")
+"""The target data type."""
+
+
+class TextSample(NamedTuple, Generic[TargetType]):
+    """Text sample with target and metadata."""
+
+    text: MessageSeries
+    """One or multiple conversation messages."""
+
+    target: TargetType | None
+    """Target data."""
+
+    metadata: dict[str, Any] | None
+    """Additional metadata."""
+
+
+class PredictionSample(NamedTuple, Generic[TargetType]):
+    """Text sample with target and metadata."""
+
+    prediction: TargetType
+    """Prediction data."""
+
+    target: TargetType
+    """Target data."""
+
+    text: MessageSeries | None
+    """Conversation messages that were used as input."""
+
+    metadata: dict[str, Any] | None
+    """Additional metadata."""

eva/language/data/messages.py

@@ -0,0 +1,60 @@
+"""Types and classes for conversation messages in a multimodal context."""
+
+import dataclasses
+import enum
+from typing import Any, Dict, List
+
+
+class Role(str, enum.Enum):
+    """Roles for messages in a conversation."""
+
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+
+
+@dataclasses.dataclass
+class Message:
+    """Base class for a message in a conversation."""
+
+    content: str
+    role: str
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the message to a dictionary."""
+        return dataclasses.asdict(self)
+
+
+@dataclasses.dataclass
+class UserMessage(Message):
+    """User message in a conversation."""
+
+    role: str = Role.USER
+
+
+@dataclasses.dataclass
+class AssistantMessage(Message):
+    """Assistant message in a conversation."""
+
+    role: str = Role.ASSISTANT
+
+
+@dataclasses.dataclass
+class SystemMessage(Message):
+    """System message in a conversation."""
+
+    role: str = Role.SYSTEM
+
+
+@dataclasses.dataclass
+class ModelSystemMessage(SystemMessage):
+    """System message for model-specific instructions."""
+
+
+@dataclasses.dataclass
+class TaskSystemMessage(SystemMessage):
+    """System message for task-specific instructions."""
+
+
+MessageSeries = List[Message]
+"""A series of conversation messages, can contain a mix of system, user, and AI messages."""

eva/language/models/__init__.py

@@ -1,25 +1,29 @@
 """Language Models API."""
 
-from eva.language.models import modules, wrappers
-from eva.language.models.modules import TextModule
-from eva.language.models.wrappers import HuggingFaceTextModel, LiteLLMTextModel
+from eva.language.models import modules, networks, wrappers
+from eva.language.models.modules import LanguageModule, OfflineLanguageModule
+from eva.language.models.wrappers import HuggingFaceModel, LiteLLMModel
 
 try:
-    from eva.language.models.wrappers import VLLMTextModel
+    from eva.language.models.wrappers import VllmModel
 
     __all__ = [
         "modules",
         "wrappers",
-        "TextModule",
-        "HuggingFaceTextModel",
-        "LiteLLMTextModel",
-        "VLLMTextModel",
+        "networks",
+        "HuggingFaceModel",
+        "LiteLLMModel",
+        "VllmModel",
+        "LanguageModule",
+        "OfflineLanguageModule",
     ]
 except ImportError:
     __all__ = [
         "modules",
         "wrappers",
-        "TextModule",
-        "HuggingFaceTextModel",
-        "LiteLLMTextModel",
+        "networks",
+        "HuggingFaceModel",
+        "LiteLLMModel",
+        "LanguageModule",
+        "OfflineLanguageModule",
     ]

eva/language/models/modules/__init__.py

@@ -1,5 +1,5 @@
 """Language Networks API."""
 
-from eva.language.models.modules.text import TextModule
+from eva.language.models.modules.language import LanguageModule, OfflineLanguageModule
 
-__all__ = ["TextModule"]
+__all__ = ["LanguageModule", "OfflineLanguageModule"]

eva/language/models/modules/language.py

@@ -0,0 +1,93 @@
+"""Model module for 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.language.models.typings import PredictionBatch, TextBatch
+
+
+class LanguageModule(module.ModelModule):
+    """Model module for 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: TextBatch, *args: Any, **kwargs: Any) -> List[str]:
+        return self.model(batch)
+
+    @override
+    def validation_step(self, batch: TextBatch, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    @override
+    def test_step(self, batch: TextBatch, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    def _batch_step(self, batch: TextBatch) -> STEP_OUTPUT:
+        text, targets, metadata = TextBatch(*batch)
+        predictions = self.forward(batch)
+        return {
+            "inputs": text,
+            "predictions": predictions,
+            "targets": targets,
+            "metadata": metadata,
+        }
+
+
+class OfflineLanguageModule(module.ModelModule):
+    """Model module for offline language tasks."""
+
+    def __init__(
+        self,
+        metrics: metrics_lib.MetricsSchema | None = None,
+        postprocess: batch_postprocess.BatchPostProcess | None = None,
+    ) -> None:
+        """Initializes the text inference module.
+
+        Args:
+            metrics: Metrics schema for evaluation.
+            postprocess: A helper function to post-process model outputs before evaluation.
+        """
+        super().__init__(metrics=metrics, postprocess=postprocess)
+
+    @override
+    def forward(self, batch: PredictionBatch, *args: Any, **kwargs: Any) -> PredictionBatch:
+        return batch
+
+    @override
+    def validation_step(self, batch: PredictionBatch, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    @override
+    def test_step(self, batch: PredictionBatch, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    def _batch_step(self, batch: PredictionBatch) -> STEP_OUTPUT:
+        predictions, targets, text, metadata = PredictionBatch(*batch)
+        return {
+            "inputs": text,
+            "predictions": predictions,
+            "targets": targets,
+            "metadata": metadata,
+        }

eva/language/models/networks/__init__.py

@@ -0,0 +1,12 @@
+"""Language networks API."""
+
+from eva.language.models.networks.alibaba import Qwen205BInstruct
+from eva.language.models.networks.api import Claude35Sonnet20240620, Claude37Sonnet20250219
+from eva.language.models.networks.registry import model_registry
+
+__all__ = [
+    "Claude35Sonnet20240620",
+    "Claude37Sonnet20250219",
+    "Qwen205BInstruct",
+    "model_registry",
+]

eva/language/models/networks/alibaba.py

@@ -0,0 +1,26 @@
+"""Models from Alibaba."""
+
+import torch
+
+from eva.language.models import wrappers
+from eva.language.models.networks.registry import model_registry
+
+
+@model_registry.register("alibaba/qwen2-0-5b-instruct")
+class Qwen205BInstruct(wrappers.HuggingFaceModel):
+    """Qwen2 0.5B Instruct model."""
+
+    def __init__(self, system_prompt: str | None = None, cache_dir: str | None = None):
+        """Initialize the model."""
+        super().__init__(
+            model_name_or_path="Qwen/Qwen2-0.5B-Instruct",
+            model_kwargs={
+                "torch_dtype": torch.bfloat16,
+                "cache_dir": cache_dir,
+            },
+            generation_kwargs={
+                "max_new_tokens": 512,
+            },
+            system_prompt=system_prompt,
+            chat_mode=True,
+        )

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

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

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

@@ -0,0 +1,34 @@
+"""Models from Anthropic."""
+
+import os
+
+from eva.language.models import wrappers
+from eva.language.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)

eva/language/models/networks/registry.py

@@ -0,0 +1,5 @@
+"""Language Model Registry."""
+
+from eva.core.utils.registry import Registry
+
+model_registry = Registry()