kaiko-eva 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

eva/core/utils/factory.py

@@ -0,0 +1,66 @@
+"""Factory classes."""
+
+import inspect
+from typing import Any, Dict, Generic, Type, TypeVar
+
+from torch import nn
+from typing_extensions import override
+
+from eva.core.utils.registry import Registry, RegistryItem
+
+T = TypeVar("T")
+
+
+class Factory(Generic[T]):
+    """A base factory class for instantiating registry items of a specific type."""
+
+    def __new__(cls, registry: Registry, name: str, init_args: dict, expected_type: Type[T]) -> T:
+        """Creates the appropriate instance based on registry entry.
+
+        Args:
+            registry: The registry containing the items to instantiate.
+            name: Name of the registry item to instantiate.
+            init_args: The arguments to pass to the constructor of the registry item.
+            expected_type: The expected type of the instantiated object.
+        """
+        if name not in registry.entries():
+            raise ValueError(
+                f"Invalid name: {name}. Please choose one "
+                f"of the following: {registry.entries()}"
+            )
+
+        registry_item = registry.get(name)
+        filtered_kwargs = _filter_kwargs(registry_item, init_args)
+        instance = registry_item(**filtered_kwargs)
+
+        if not isinstance(instance, expected_type):
+            raise TypeError(f"Expected an instance of {expected_type}, but got {type(instance)}.")
+        return instance
+
+
+class ModuleFactory(Factory[nn.Module]):
+    """Factory class for instantiating nn.Module instances from a registry."""
+
+    @override
+    def __new__(cls, registry: Registry, name: str, init_args: dict) -> nn.Module:
+        return super().__new__(cls, registry, name, init_args, nn.Module)
+
+
+def _filter_kwargs(registry_item: RegistryItem, kwargs: dict) -> Dict[str, Any]:
+    """Filters the given keyword arguments to match the signature of a given class or method.
+
+    Args:
+        registry_item: The class or method from the registry whose
+            signature should be used for filtering.
+        kwargs: The keyword arguments to filter.
+
+    Returns:
+        A dictionary containing only the valid keyword arguments that match
+        the callable's parameters.
+    """
+    if inspect.isclass(registry_item):
+        signature = inspect.signature(registry_item.__init__)
+    else:
+        signature = inspect.signature(registry_item)
+
+    return {k: v for k, v in kwargs.items() if k in signature.parameters}

eva/core/utils/registry.py

@@ -0,0 +1,42 @@
+"""Registry for classes and methods."""
+
+from typing import Any, Callable, Dict, List, Type, Union
+
+RegistryItem = Union[Type[Any], Callable[..., Any]]
+
+
+class Registry:
+    """A registry to store and access classes and methods by a unique key."""
+
+    def __init__(self) -> None:
+        """Initializes the registry class."""
+        self._registry: Dict[str, RegistryItem] = {}
+
+    def register(self, key: str, /) -> Callable[[RegistryItem], RegistryItem]:
+        """A decorator to register a class or method with a unique key.
+
+        Args:
+            key: The key to register the class or method under.
+
+        Returns:
+            A decorator that registers the class or method in the registry.
+        """
+
+        def wrapper(obj: RegistryItem) -> RegistryItem:
+            if key in self.entries():
+                raise ValueError(f"Entry {key} is already registered.")
+
+            self._registry[key] = obj
+            return obj
+
+        return wrapper
+
+    def get(self, name: str) -> RegistryItem:
+        """Gets the class or method from the registry."""
+        if name not in self._registry:
+            raise ValueError(f"Item {name} not found in the registry.")
+        return self._registry[name]
+
+    def entries(self) -> List[str]:
+        """List all items in the registry."""
+        return list(self._registry.keys())

eva/core/utils/requirements.py

@@ -0,0 +1,26 @@
+"""Utility functions related to package requirements."""
+
+import importlib
+from typing import Dict
+
+from packaging import version
+
+
+def check_dependencies(requirements: Dict[str, str]) -> None:
+    """Check installed package versions against requirements dict.
+
+    Args:
+        requirements: A dictionary where keys are package names and
+            values are minimum required versions.
+
+    Raises:
+        ImportError: If any package does not meet the minimum required version.
+    """
+    for package, min_version in requirements.items():
+        module = importlib.import_module(package)
+        actual = getattr(module, "__version__", None)
+        if actual and not (version.parse(actual) >= version.parse(min_version)):
+            raise ImportError(
+                f"Package '{package}' version {actual} does not meet "
+                f"the minimum required version {min_version}."
+            )

eva/language/__init__.py

@@ -0,0 +1,13 @@
+"""eva language API."""
+
+try:
+    from eva.language.data import datasets
+except ImportError as e:
+    msg = (
+        "eva language requirements are not installed.\n\n"
+        "Please pip install as follows:\n"
+        '  python -m pip install "kaiko-eva[language]" --upgrade'
+    )
+    raise ImportError(str(e) + "\n\n" + msg) from e
+
+__all__ = ["datasets"]

eva/language/data/__init__.py

@@ -0,0 +1,5 @@
+"""Language data API."""
+
+from eva.language.data import datasets
+
+__all__ = ["datasets"]

eva/language/data/datasets/__init__.py

@@ -0,0 +1,9 @@
+"""Language Datasets API."""
+
+from eva.language.data.datasets.classification import PubMedQA
+from eva.language.data.datasets.language import LanguageDataset
+
+__all__ = [
+    "PubMedQA",
+    "LanguageDataset",
+]

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

@@ -0,0 +1,7 @@
+"""Text classification datasets API."""
+
+from eva.language.data.datasets.classification.pubmedqa import PubMedQA
+
+__all__ = [
+    "PubMedQA",
+]

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

@@ -0,0 +1,63 @@
+"""Base for text classification datasets."""
+
+import abc
+from typing import Any, Dict, List, Tuple
+
+import torch
+from typing_extensions import override
+
+from eva.language.data.datasets.language import LanguageDataset
+
+
+class TextClassification(LanguageDataset[Tuple[str, torch.Tensor, Dict[str, Any]]], abc.ABC):
+    """Text classification abstract dataset."""
+
+    def __init__(self) -> None:
+        """Initializes the text classification dataset."""
+        super().__init__()
+
+    @property
+    def classes(self) -> List[str] | None:
+        """Returns list of class names."""
+
+    @property
+    def class_to_idx(self) -> Dict[str, int] | None:
+        """Returns class name to index mapping."""
+
+    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) -> str:
+        """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) -> torch.Tensor:
+        """Returns the target label.
+
+        Args:
+            index: The index of the data sample.
+
+        Returns:
+            The target label.
+        """
+        raise NotImplementedError
+
+    @override
+    def __getitem__(self, index: int) -> Tuple[str, torch.Tensor, Dict[str, Any]]:
+        return (self.load_text(index), self.load_target(index), self.load_metadata(index) or {})

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

@@ -0,0 +1,149 @@
+"""PubMedQA dataset class."""
+
+import os
+import random
+from typing import Dict, List, Literal
+
+import torch
+from datasets import Dataset, load_dataset, load_from_disk
+from loguru import logger
+from typing_extensions import override
+
+from eva.language.data.datasets.classification import base
+
+
+class PubMedQA(base.TextClassification):
+    """Dataset class for PubMedQA question answering task."""
+
+    _license: str = "MIT License (https://github.com/pubmedqa/pubmedqa/blob/master/LICENSE)"
+    """Dataset license."""
+
+    def __init__(
+        self,
+        root: str | None = None,
+        split: Literal["train", "val", "test"] | None = None,
+        download: bool = False,
+        max_samples: int | None = None,
+    ) -> None:
+        """Initialize the PubMedQA dataset.
+
+        Args:
+            root: Directory to cache the dataset. If None, no local caching is used.
+            split: Valid splits among ["train", "val", "test"].
+                If None, it will use "train+test+validation".
+            download: Whether to download the dataset if not found locally. Default is False.
+            max_samples: Maximum number of samples to use. If None, use all samples.
+        """
+        super().__init__()
+
+        self._root = root
+        self._split = split
+        self._download = download
+        self._max_samples = max_samples
+
+    def _load_dataset(self, dataset_path: str | None) -> Dataset:
+        """Loads the PubMedQA dataset from the local cache or downloads it.
+
+        Args:
+            dataset_path: The path to the local cache (may be None).
+
+        Returns:
+            The loaded dataset object.
+        """
+        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"
+
+        if self._download:
+            logger.info("Downloading dataset from HuggingFace Hub")
+            raw_dataset = load_dataset(
+                dataset_name,
+                name=config_name,
+                split=split,
+                trust_remote_code=True,
+                download_mode="reuse_dataset_if_exists",
+            )
+            if dataset_path:
+                raw_dataset.save_to_disk(dataset_path)  # type: ignore
+                logger.info(f"Dataset saved to: {dataset_path}")
+        else:
+            if not dataset_path or not os.path.exists(dataset_path):
+                raise ValueError(
+                    "Dataset path not found. Set download=True or provide a valid root path."
+                )
+
+            logger.info(f"Loading dataset from: {dataset_path}")
+            raw_dataset = load_from_disk(dataset_path)
+
+        return raw_dataset  # type: ignore
+
+    @override
+    def prepare_data(self) -> None:
+        """Downloads and prepares the PubMedQA dataset.
+
+        If `self._root` is None, the dataset is used directly from HuggingFace.
+        Otherwise, it checks if the dataset is already cached in `self._root`.
+        If not cached, it downloads the dataset into `self._root`.
+        """
+        dataset_path = None
+
+        if self._root:
+            dataset_path = self._root
+            os.makedirs(self._root, exist_ok=True)
+
+        try:
+            self.dataset = self._load_dataset(dataset_path)
+            if self._max_samples is not None and len(self.dataset) > self._max_samples:
+                logger.info(
+                    f"Subsampling dataset from {len(self.dataset)} to {self._max_samples} samples"
+                )
+                random.seed(42)
+                indices = random.sample(range(len(self.dataset)), self._max_samples)
+                self.dataset = self.dataset.select(indices)
+        except Exception as e:
+            raise RuntimeError(f"Failed to prepare dataset: {e}") from e
+
+    @property
+    @override
+    def classes(self) -> List[str]:
+        return ["no", "yes", "maybe"]
+
+    @property
+    @override
+    def class_to_idx(self) -> Dict[str, int]:
+        return {"no": 0, "yes": 1, "maybe": 2}
+
+    @override
+    def load_text(self, index: int) -> str:
+        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"])
+
+    @override
+    def load_target(self, index: int) -> torch.Tensor:
+        if index < 0 or index >= len(self.dataset):
+            raise IndexError(f"Index {index} out of range for dataset of size {len(self.dataset)}")
+        return torch.tensor(
+            self.class_to_idx[self.dataset[index]["final_decision"]], dtype=torch.long
+        )
+
+    @override
+    def load_metadata(self, index: int) -> Dict[str, str]:
+        sample = self.dataset[index]
+        return {
+            "year": sample.get("YEAR") or "",
+            "labels": sample.get("LABELS") or "",
+            "meshes": sample.get("MESHES") or "",
+            "long_answer": sample.get("LONG_ANSWER") or "",
+            "reasoning_required": sample.get("reasoning_required_pred") or "",
+            "reasoning_free": sample.get("reasoning_free_pred") or "",
+        }
+
+    @override
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def _print_license(self) -> None:
+        """Prints the dataset license."""
+        print(f"Dataset license: {self._license}")

eva/language/data/datasets/language.py

@@ -0,0 +1,13 @@
+"""Language 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 LanguageDataset(base.MapDataset, abc.ABC, Generic[DataSample]):
+    """Base dataset class for text tasks."""

eva/language/models/__init__.py

@@ -0,0 +1,25 @@
+"""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
+
+try:
+    from eva.language.models.wrappers import VLLMTextModel
+
+    __all__ = [
+        "modules",
+        "wrappers",
+        "TextModule",
+        "HuggingFaceTextModel",
+        "LiteLLMTextModel",
+        "VLLMTextModel",
+    ]
+except ImportError:
+    __all__ = [
+        "modules",
+        "wrappers",
+        "TextModule",
+        "HuggingFaceTextModel",
+        "LiteLLMTextModel",
+    ]

eva/language/models/modules/__init__.py

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

eva/language/models/modules/text.py

@@ -0,0 +1,85 @@
+"""LLM Text Module for Inference."""
+
+from typing import Any, List
+
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from loguru import logger
+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.modules.typings import TEXT_BATCH
+
+
+class TextModule(module.ModelModule):
+    """Text-based LLM module for inference.
+
+    Uses LLM wrappers for text generation and supports evaluation using
+    configurable metrics and post-processing transforms.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        prompt: str,
+        metrics: metrics_lib.MetricsSchema | None = None,
+        postprocess: batch_postprocess.BatchPostProcess | None = None,
+    ) -> None:
+        """Initializes the text inference module.
+
+        Args:
+            model: An LLM wrapper (PyTorch-compatible) for text generation.
+            prompt: The prompt to use for generating text.
+            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
+        self.prompt = prompt
+
+    @override
+    def forward(self, prompts: List[str], *args: Any, **kwargs: Any) -> List[str]:
+        """Generates text responses for a batch of prompts.
+
+        Args:
+            prompts: List of input texts to generate responses.
+            args: Additional arguments.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            List of generated responses.
+        """
+        return self.model(prompts)
+
+    @override
+    def validation_step(self, batch: TEXT_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        """Validation step that runs batch inference and evaluates metrics.
+
+        Args:
+            batch: An input batch.
+            args: Additional arguments.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            Dictionary with predictions, ground truth, and evaluation metrics.
+        """
+        return self._batch_step(batch)
+
+    def _batch_step(self, batch: TEXT_BATCH) -> STEP_OUTPUT:
+        """Runs inference on a batch and evaluates model predictions.
+
+        Args:
+            batch: Input batch containing data, targets, and metadata.
+
+        Returns:
+            Dictionary with predictions, ground truth, and evaluation metrics.
+        """
+        data, targets, metadata = batch
+        messages = [str(d) + "\n" + self.prompt for d in data]
+        predictions = self(messages)
+        logger.debug(f"Predictions: {predictions}")
+        logger.debug(f"Targets: {targets}")
+        return {"predictions": predictions, "targets": targets, "metadata": metadata}

eva/language/models/modules/typings.py

@@ -0,0 +1,16 @@
+"""Type annotations for language model modules."""
+
+from typing import Any, Dict, List, NamedTuple
+
+
+class TEXT_BATCH(NamedTuple):
+    """Text-based input batch data scheme."""
+
+    data: List[str]
+    """The text data batch."""
+
+    targets: List[str] | None = None
+    """The target text batch."""
+
+    metadata: Dict[str, Any] | None = None
+    """The associated metadata."""

eva/language/models/wrappers/__init__.py

@@ -0,0 +1,11 @@
+"""Language Model Wrappers API."""
+
+from eva.language.models.wrappers.huggingface import HuggingFaceTextModel
+from eva.language.models.wrappers.litellm import LiteLLMTextModel
+
+try:
+    from eva.language.models.wrappers.vllm import VLLMTextModel
+
+    __all__ = ["HuggingFaceTextModel", "LiteLLMTextModel", "VLLMTextModel"]
+except ImportError:
+    __all__ = ["HuggingFaceTextModel", "LiteLLMTextModel"]

eva/language/models/wrappers/huggingface.py

@@ -0,0 +1,69 @@
+"""LLM wrapper for HuggingFace `transformers` models."""
+
+from typing import Any, Dict, List, Literal
+
+from transformers.pipelines import pipeline
+from typing_extensions import override
+
+from eva.core.models.wrappers import base
+
+
+class HuggingFaceTextModel(base.BaseModel[List[str], List[str]]):
+    """Wrapper class for loading HuggingFace `transformers` models using pipelines."""
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        task: Literal["text-generation"] = "text-generation",
+        model_kwargs: Dict[str, Any] | None = None,
+        generation_kwargs: Dict[str, Any] | None = None,
+    ) -> None:
+        """Initializes the model.
+
+        Args:
+            model_name_or_path: The model name or path to load the model from.
+                This can be a local path or a model name from the `HuggingFace`
+                model hub.
+            task: The pipeline task. Defaults to "text-generation".
+            model_kwargs: Additional arguments for configuring the pipeline.
+            generation_kwargs: Additional generation parameters (temperature, max_length, etc.).
+        """
+        super().__init__()
+
+        self._model_name_or_path = model_name_or_path
+        self._task = task
+        self._model_kwargs = model_kwargs or {}
+        self._generation_kwargs = generation_kwargs or {}
+
+        self.load_model()
+
+    @override
+    def load_model(self) -> None:
+        """Loads the model as a Hugging Face pipeline."""
+        self._pipeline = pipeline(
+            task=self._task,
+            model=self._model_name_or_path,
+            trust_remote_code=True,
+            **self._model_kwargs,
+        )
+
+    @override
+    def model_forward(self, prompts: List[str]) -> List[str]:
+        """Generates text using the pipeline.
+
+        Args:
+            prompts: The input prompts for the model.
+
+        Returns:
+            The generated text as a string.
+        """
+        outputs = self._pipeline(prompts, return_full_text=False, **self._generation_kwargs)
+        if outputs is None:
+            raise ValueError("Outputs from the model are None.")
+        results = []
+        for output in outputs:
+            if isinstance(output, list):
+                results.append(output[0]["generated_text"])  # type: ignore
+            else:
+                results.append(output["generated_text"])  # type: ignore
+        return results