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

eva/core/data/dataloaders/__init__.py

@@ -1,5 +1,6 @@
 """Dataloaders API."""
 
+from eva.core.data.dataloaders.collate_fn import text_collate_fn
 from eva.core.data.dataloaders.dataloader import DataLoader
 
-__all__ = ["DataLoader"]
+__all__ = ["text_collate_fn", "DataLoader"]

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

@@ -0,0 +1,5 @@
+"""Collate functions API."""
+
+from eva.core.data.dataloaders.collate_fn.collate import text_collate_fn
+
+__all__ = ["text_collate_fn"]

eva/core/data/dataloaders/collate_fn/collate.py

@@ -0,0 +1,24 @@
+"""Collate functions for text data."""
+
+from typing import Dict, List, Tuple
+
+import torch
+
+
+def text_collate_fn(
+    batch: List[Tuple[str, torch.Tensor, Dict]],
+) -> Tuple[List[str], torch.Tensor, List[Dict]]:
+    """Collate function for text data that keeps texts as separate strings.
+
+    Args:
+        batch: List of tuples containing (text, target, metadata) from the dataset
+
+    Returns:
+        Tuple containing:
+            - List of text strings
+            - Batched tensor of targets
+            - List of metadata dictionaries
+    """
+    texts, targets, metadata = zip(*batch, strict=False)
+    targets = torch.stack(targets)
+    return list(texts), targets, list(metadata)

eva/core/data/dataloaders/dataloader.py

@@ -56,6 +56,9 @@ class DataLoader:
     persistent_workers: bool = True
     """Will keep the worker processes after a dataset has been consumed once."""
 
+    worker_init_fn: Callable | None = None
+    """Function to call on each worker process before data loading."""
+
     prefetch_factor: int | None = 2
     """Number of batches loaded in advance by each worker."""
 
@@ -80,4 +83,5 @@ class DataLoader:
             drop_last=self.drop_last,
             persistent_workers=self.persistent_workers,
             prefetch_factor=self.prefetch_factor,
+            worker_init_fn=self.worker_init_fn,
         )

eva/core/interface/interface.py

@@ -34,7 +34,14 @@ class Interface:
             model: The model module to use but not modify.
             data: The data module.
         """
-        trainer.run_evaluation_session(model=model, datamodule=data)
+        eva_trainer.run_evaluation_session(
+            base_trainer=trainer,
+            base_model=model,
+            datamodule=data,
+            stages=["fit", "validate", "test"],
+            n_runs=trainer.n_runs,
+            verbose=trainer.n_runs > 1,
+        )
 
     def predict(
         self,
@@ -77,3 +84,29 @@ class Interface:
         """
         self.predict(trainer=trainer, model=model, data=data)
         self.fit(trainer=trainer, model=model, data=data)
+
+    def validate(
+        self,
+        trainer: eva_trainer.Trainer,
+        model: modules.ModelModule,
+        data: datamodules.DataModule,
+    ) -> None:
+        """Perform model validation out-of-place without running fit.
+
+        This method is useful when the model is already trained or does not
+        require further training (e.g., large language models) and you only
+        want to measure performance.
+
+        Args:
+            trainer: The base trainer to use but not modify.
+            model: The model module to use but not modify.
+            data: The data module containing validation data.
+        """
+        eva_trainer.run_evaluation_session(
+            base_trainer=trainer,
+            base_model=model,
+            datamodule=data,
+            stages=["validate"],
+            n_runs=trainer.n_runs,
+            verbose=trainer.n_runs > 1,
+        )

eva/core/metrics/defaults/classification/multiclass.py

@@ -17,6 +17,7 @@ class MulticlassClassificationMetrics(structs.MetricCollection):
         ignore_index: int | None = None,
         prefix: str | None = None,
         postfix: str | None = None,
+        input_type: Literal["logits", "discrete"] = "logits",
     ) -> None:
         """Initializes the multi-class classification metrics.
 
@@ -27,46 +28,55 @@ class MulticlassClassificationMetrics(structs.MetricCollection):
                 contribute to the metric calculation.
             prefix: A string to append in front of the keys of the output dict.
             postfix: A string to append after the keys of the output dict.
+            input_type: Type of input predictions - "logits" for probabilities/logits
+                or "discrete" for discrete class predictions. Determines which metrics
+                are applicable.
         """
-        super().__init__(
-            metrics=[
+        metrics = [
+            classification.MulticlassAccuracy(
+                num_classes=num_classes,
+                average=average,
+                ignore_index=ignore_index,
+            ),
+            classification.MulticlassF1Score(
+                num_classes=num_classes,
+                average=average,
+                ignore_index=ignore_index,
+            ),
+            classification.MulticlassPrecision(
+                num_classes=num_classes,
+                average=average,
+                ignore_index=ignore_index,
+            ),
+            classification.MulticlassRecall(
+                num_classes=num_classes,
+                average=average,
+                ignore_index=ignore_index,
+            ),
+        ]
+
+        compute_groups = [
+            [
+                "MulticlassAccuracy",
+                "MulticlassF1Score",
+                "MulticlassPrecision",
+                "MulticlassRecall",
+            ]
+        ]
+
+        if input_type == "logits":
+            metrics.append(
                 classification.MulticlassAUROC(
                     num_classes=num_classes,
                     average=average,
                     ignore_index=ignore_index,
-                ),
-                classification.MulticlassAccuracy(
-                    num_classes=num_classes,
-                    average=average,
-                    ignore_index=ignore_index,
-                ),
-                classification.MulticlassF1Score(
-                    num_classes=num_classes,
-                    average=average,
-                    ignore_index=ignore_index,
-                ),
-                classification.MulticlassPrecision(
-                    num_classes=num_classes,
-                    average=average,
-                    ignore_index=ignore_index,
-                ),
-                classification.MulticlassRecall(
-                    num_classes=num_classes,
-                    average=average,
-                    ignore_index=ignore_index,
-                ),
-            ],
+                )
+            )
+            compute_groups.append(["MulticlassAUROC"])
+
+        super().__init__(
+            metrics=metrics,
             prefix=prefix,
             postfix=postfix,
-            compute_groups=[
-                [
-                    "MulticlassAccuracy",
-                    "MulticlassF1Score",
-                    "MulticlassPrecision",
-                    "MulticlassRecall",
-                ],
-                [
-                    "MulticlassAUROC",
-                ],
-            ],
+            compute_groups=compute_groups,
         )

eva/core/models/modules/__init__.py

@@ -3,5 +3,6 @@
 from eva.core.models.modules.head import HeadModule
 from eva.core.models.modules.inference import InferenceModule
 from eva.core.models.modules.module import ModelModule
+from eva.core.models.modules.scheduler import SchedulerConfiguration
 
-__all__ = ["HeadModule", "ModelModule", "InferenceModule"]
+__all__ = ["HeadModule", "ModelModule", "InferenceModule", "SchedulerConfiguration"]

eva/core/models/modules/scheduler.py

@@ -0,0 +1,51 @@
+"""Learning Rate scheduler configuration."""
+
+import dataclasses
+from typing import Any, Literal
+
+from lightning.pytorch.cli import LRSchedulerCallable
+from torch import optim
+
+
+@dataclasses.dataclass
+class SchedulerConfiguration:
+    """Initializes and builds the learning rate scheduler configuration."""
+
+    scheduler: LRSchedulerCallable
+    """The learning rate scheduler instance."""
+
+    interval: Literal["step", "epoch"] = "epoch"
+    """The unit of the scheduler's step size.
+
+    It can be 'step' or 'epoch', to update the scheduler on step or epoch end respectively.
+    """
+
+    frequency: int = 1
+    """How many epochs/steps should pass between calls to `scheduler.step()`.
+
+    Value `1` corresponds to updating the learning rate after every epoch/step.
+    """
+
+    monitor: str = "val_loss"
+    """Metric to to monitor for schedulers like `ReduceLROnPlateau`."""
+
+    strict: bool = True
+    """Whether to enforce that the value specified 'monitor' must be available.
+
+    If the values is not available when the scheduler is updated it will stop the
+    training. With `False`, it will only produce a warning.
+    """
+
+    name: str | None = None
+    """Specifies a custom logged name for the `LearningRateMonitor` callback."""
+
+    def __call__(self, optimizer: optim.Optimizer) -> dict[str, Any]:
+        """Returns Lightning's lr_scheduler_config configuration."""
+        return {
+            "scheduler": self.scheduler(optimizer),
+            "interval": self.interval,
+            "frequency": self.frequency,
+            "monitor": self.monitor,
+            "strict": self.strict,
+            "name": self.name,
+        }

eva/core/models/transforms/extract_cls_features.py

@@ -31,7 +31,7 @@ class ExtractCLSFeatures:
             tensor: The tensor representing the model output.
         """
         if isinstance(tensor, modeling_outputs.BaseModelOutputWithPooling):
-            tensor = tensor.last_hidden_state
+            tensor = tensor.last_hidden_state  # type: ignore
 
         cls_token = tensor[:, self._cls_index, :]
         if self._include_patch_tokens:

eva/core/models/transforms/extract_patch_features.py

@@ -43,7 +43,7 @@ class ExtractPatchFeatures:
         """
         num_skip = int(self._has_cls_token) + self._num_register_tokens
         if isinstance(tensor, modeling_outputs.BaseModelOutputWithPooling):
-            features = tensor.last_hidden_state[:, num_skip:, :].permute(0, 2, 1)
+            features = tensor.last_hidden_state[:, num_skip:, :].permute(0, 2, 1)  # type: ignore
         else:
             features = tensor[:, num_skip:, :].permute(0, 2, 1)
 

eva/core/models/wrappers/base.py

@@ -1,40 +1,43 @@
 """Base class for model wrappers."""
 
 import abc
-from typing import Callable
+from typing import Callable, Generic, TypeVar
 
-import torch
 import torch.nn as nn
 from typing_extensions import override
 
+InputType = TypeVar("InputType")
+"""The input data type."""
+OutputType = TypeVar("OutputType")
+"""The output data type."""
 
-class BaseModel(nn.Module):
+
+class BaseModel(nn.Module, Generic[InputType, OutputType]):
     """Base class for model wrappers."""
 
-    def __init__(self, tensor_transforms: Callable | None = None) -> None:
+    def __init__(self, transforms: Callable | None = None) -> None:
         """Initializes the model.
 
         Args:
-            tensor_transforms: The transforms to apply to the output
-                tensor produced by the model.
+            transforms: The transforms to apply to the output produced by the model.
         """
         super().__init__()
 
-        self._output_transforms = tensor_transforms
+        self._output_transforms = transforms
 
-        self._model: Callable[..., torch.Tensor] | nn.Module
+        self._model: Callable[..., OutputType] | nn.Module
 
     @override
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        tensor = self.model_forward(tensor)
-        return self._apply_transforms(tensor)
+    def forward(self, tensor: InputType) -> OutputType:
+        out = self.model_forward(tensor)
+        return self._apply_transforms(out)
 
     @abc.abstractmethod
-    def load_model(self) -> Callable[..., torch.Tensor]:
+    def load_model(self) -> Callable[..., OutputType]:
         """Loads the model."""
         raise NotImplementedError
 
-    def model_forward(self, tensor: torch.Tensor) -> torch.Tensor:
+    def model_forward(self, tensor: InputType) -> OutputType:
         """Implements the forward pass of the model.
 
         Args:
@@ -42,7 +45,7 @@ class BaseModel(nn.Module):
         """
         return self._model(tensor)
 
-    def _apply_transforms(self, tensor: torch.Tensor) -> torch.Tensor:
+    def _apply_transforms(self, tensor: OutputType) -> OutputType:
         if self._output_transforms is not None:
             tensor = self._output_transforms(tensor)
         return tensor

eva/core/models/wrappers/from_function.py

@@ -3,13 +3,14 @@
 from typing import Any, Callable, Dict
 
 import jsonargparse
+import torch
 from torch import nn
 from typing_extensions import override
 
 from eva.core.models.wrappers import _utils, base
 
 
-class ModelFromFunction(base.BaseModel):
+class ModelFromFunction(base.BaseModel[torch.Tensor, torch.Tensor]):
     """Wrapper class for models which are initialized from functions.
 
     This is helpful for initializing models in a `.yaml` configuration file.
@@ -20,7 +21,7 @@ class ModelFromFunction(base.BaseModel):
         path: Callable[..., nn.Module],
         arguments: Dict[str, Any] | None = None,
         checkpoint_path: str | None = None,
-        tensor_transforms: Callable | None = None,
+        transforms: Callable | None = None,
     ) -> None:
         """Initializes and constructs the model.
 
@@ -31,10 +32,10 @@ class ModelFromFunction(base.BaseModel):
                 weights from. This is currently only supported for torch
                 model checkpoints. For other formats, the checkpoint loading
                 should be handled within the provided callable object in <path>.
-            tensor_transforms: The transforms to apply to the output tensor
+            transforms: The transforms to apply to the output tensor
                 produced by the model.
         """
-        super().__init__(tensor_transforms=tensor_transforms)
+        super().__init__(transforms=transforms)
 
         self._path = path
         self._arguments = arguments

eva/core/models/wrappers/from_torchhub.py

@@ -6,11 +6,10 @@ import torch
 import torch.nn as nn
 from typing_extensions import override
 
-from eva.core.models import wrappers
-from eva.core.models.wrappers import _utils
+from eva.core.models.wrappers import _utils, base
 
 
-class TorchHubModel(wrappers.BaseModel):
+class TorchHubModel(base.BaseModel[torch.Tensor, torch.Tensor]):
     """Model wrapper for `torch.hub` models."""
 
     def __init__(
@@ -23,7 +22,7 @@ class TorchHubModel(wrappers.BaseModel):
         norm: bool = False,
         trust_repo: bool = True,
         model_kwargs: Dict[str, Any] | None = None,
-        tensor_transforms: Callable | None = None,
+        transforms: Callable | None = None,
     ) -> None:
         """Initializes the encoder.
 
@@ -39,10 +38,10 @@ class TorchHubModel(wrappers.BaseModel):
             trust_repo: If set to `False`, a prompt will ask the user whether the
                 repo should be trusted.
             model_kwargs: Extra model arguments.
-            tensor_transforms: The transforms to apply to the output tensor
+            transforms: The transforms to apply to the output tensor
                 produced by the model.
         """
-        super().__init__(tensor_transforms=tensor_transforms)
+        super().__init__(transforms=transforms)
 
         self._model_name = model_name
         self._repo_or_dir = repo_or_dir

eva/core/models/wrappers/huggingface.py

@@ -2,19 +2,20 @@
 
 from typing import Any, Callable, Dict
 
+import torch
 import transformers
 from typing_extensions import override
 
 from eva.core.models.wrappers import base
 
 
-class HuggingFaceModel(base.BaseModel):
+class HuggingFaceModel(base.BaseModel[torch.Tensor, torch.Tensor]):
     """Wrapper class for loading HuggingFace `transformers` models."""
 
     def __init__(
         self,
         model_name_or_path: str,
-        tensor_transforms: Callable | None = None,
+        transforms: Callable | None = None,
         model_kwargs: Dict[str, Any] | None = None,
     ) -> None:
         """Initializes the model.
@@ -23,11 +24,11 @@ class HuggingFaceModel(base.BaseModel):
             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.
-            tensor_transforms: The transforms to apply to the output tensor
+            transforms: The transforms to apply to the output tensor
                 produced by the model.
             model_kwargs: The arguments used for instantiating the model.
         """
-        super().__init__(tensor_transforms=tensor_transforms)
+        super().__init__(transforms=transforms)
 
         self._model_name_or_path = model_name_or_path
         self._model_kwargs = model_kwargs or {}
@@ -36,6 +37,8 @@ class HuggingFaceModel(base.BaseModel):
 
     @override
     def load_model(self) -> None:
+        # Use safetensors to avoid torch.load security vulnerability
+        model_kwargs = {"use_safetensors": True, **self._model_kwargs}
         self._model = transformers.AutoModel.from_pretrained(
-            self._model_name_or_path, **self._model_kwargs
+            self._model_name_or_path, **model_kwargs
         )

eva/core/models/wrappers/onnx.py

@@ -9,23 +9,23 @@ from typing_extensions import override
 from eva.core.models.wrappers import base
 
 
-class ONNXModel(base.BaseModel):
+class ONNXModel(base.BaseModel[torch.Tensor, torch.Tensor]):
     """Wrapper class for loading ONNX models."""
 
     def __init__(
         self,
         path: str,
         device: Literal["cpu", "cuda"] | None = "cpu",
-        tensor_transforms: Callable | None = None,
+        transforms: Callable | None = None,
     ):
         """Initializes the model.
 
         Args:
             path: The path to the .onnx model file.
             device: The device to run the model on. This can be either "cpu" or "cuda".
-            tensor_transforms: The transforms to apply to the output tensor produced by the model.
+            transforms: The transforms to apply to the output tensor produced by the model.
         """
-        super().__init__(tensor_transforms=tensor_transforms)
+        super().__init__(transforms=transforms)
 
         self._path = path
         self._device = device

eva/core/trainers/functional.py

@@ -1,6 +1,6 @@
 """Fit session related functions."""
 
-from typing import Tuple
+from typing import List, Literal, Tuple
 
 from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT
 
@@ -16,11 +16,12 @@ def run_evaluation_session(
     datamodule: datamodules.DataModule,
     *,
     n_runs: int = 1,
+    stages: List[Literal["fit", "validate", "test"]] | None = None,
     verbose: bool = True,
 ) -> None:
     """Runs a downstream evaluation session out-of-place.
 
-    It performs an evaluation run (fit and evaluate) on the model
+    It performs an evaluation run (with configurable stages) on the model
     multiple times. Note that as the input `base_trainer` and
     `base_model` would be cloned, the input object would not
     be modified.
@@ -29,10 +30,13 @@ def run_evaluation_session(
         base_trainer: The base trainer module to use.
         base_model: The base model module to use.
         datamodule: The data module.
-        n_runs: The amount of runs (fit and evaluate) to perform.
+        n_runs: The number of runs to perform.
+        stages: List of stages to execute. Options: "fit", "validate", "test".
         verbose: Whether to verbose the session metrics instead of
-            these of each individual runs and vice-versa.
+            those of each individual run and vice-versa.
     """
+    if not stages:
+        stages = ["fit", "validate", "test"]
     recorder = _recorder.SessionRecorder(output_dir=base_trainer.default_log_dir, verbose=verbose)
     for run_index in range(n_runs):
         validation_scores, test_scores = run_evaluation(
@@ -40,9 +44,11 @@ def run_evaluation_session(
             base_model,
             datamodule,
             run_id=run_index,
+            stages=stages,
             verbose=not verbose,
         )
-        recorder.update(validation_scores, test_scores)
+        if validation_scores:
+            recorder.update(validation_scores, test_scores)
     recorder.save()
 
 
@@ -52,61 +58,52 @@ def run_evaluation(
     datamodule: datamodules.DataModule,
     *,
     run_id: int | None = None,
+    stages: List[Literal["fit", "validate", "test"]] | None = None,
     verbose: bool = True,
-) -> Tuple[_EVALUATE_OUTPUT, _EVALUATE_OUTPUT | None]:
-    """Fits and evaluates a model out-of-place.
+) -> Tuple[_EVALUATE_OUTPUT | None, _EVALUATE_OUTPUT | None]:
+    """Runs the specified evaluation stages out-of-place.
 
     Args:
         base_trainer: The base trainer to use but not modify.
         base_model: The model module to use but not modify.
         datamodule: The data module.
         run_id: The run id to be appended to the output log directory.
+            If `None`, it will use the log directory of the trainer as is.
+        stages: List of stages to execute. Options: "fit", "validate", "test".
         verbose: Whether to print the validation and test metrics
             in the end of the training.
 
     Returns:
-        A tuple of with the validation and the test metrics (if exists).
+        A tuple with the validation and the test metrics (if executed).
+        If a stage is not executed, its value will be None.
     """
+    if not stages:
+        stages = ["fit", "validate", "test"]
     trainer, model = _utils.clone(base_trainer, base_model)
     model.configure_model()
 
     trainer.init_logger_run(run_id)
-    results = fit_and_validate(trainer, model, datamodule, verbose=verbose)
-    trainer.finish_logger_run(run_id)
-
-    return results
-
-
-def fit_and_validate(
-    trainer: eva_trainer.Trainer,
-    model: modules.ModelModule,
-    datamodule: datamodules.DataModule,
-    verbose: bool = True,
-) -> Tuple[_EVALUATE_OUTPUT, _EVALUATE_OUTPUT | None]:
-    """Fits and evaluates a model in-place.
-
-    If the test set is set in the datamodule, it will evaluate the model
-    on the test set as well.
 
-    Args:
-        trainer: The trainer module to use and update in-place.
-        model: The model module to use and update in-place.
-        datamodule: The data module.
-        verbose: Whether to print the validation and test metrics
-            in the end of the training.
-
-    Returns:
-        A tuple of with the validation and the test metrics (if exists).
-    """
-    trainer.fit(model, datamodule=datamodule)
-    validation_scores = trainer.validate(
-        datamodule=datamodule, verbose=verbose, ckpt_path=trainer.checkpoint_type
-    )
-    test_scores = (
-        None
-        if datamodule.datasets.test is None
-        else trainer.test(datamodule=datamodule, verbose=verbose, ckpt_path=trainer.checkpoint_type)
-    )
+    validation_scores = None
+    test_scores = None
+
+    if "fit" in stages:
+        trainer.fit(model, datamodule=datamodule)
+    if "validate" in stages:
+        validation_scores = trainer.validate(
+            model=model,
+            datamodule=datamodule,
+            verbose=verbose,
+            ckpt_path=trainer.checkpoint_type,
+        )
+    if "test" in stages and getattr(datamodule.datasets, "test", None) is not None:
+        test_scores = trainer.test(
+            model=model,
+            datamodule=datamodule,
+            verbose=verbose,
+            ckpt_path=trainer.checkpoint_type,
+        )
+    trainer.finish_logger_run(run_id)
     return validation_scores, test_scores