kaiko-eva 0.2.1__tar.gz → 0.3.0__tar.gz

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kaiko-eva
-Version: 0.2.1
+Version: 0.3.0
 Summary: Evaluation Framework for oncology foundation models.
 Keywords: machine-learning,evaluation-framework,oncology,foundation-models
 Author-Email: Ioannis Gatopoulos <ioannis@kaiko.ai>, =?utf-8?q?Nicolas_K=C3=A4nzig?= <nicolas@kaiko.ai>, Roman Moser <roman@kaiko.ai>
@@ -242,6 +242,9 @@ Requires-Dist: imagesize>=1.4.1; extra == "vision"
 Requires-Dist: scipy>=1.14.0; extra == "vision"
 Requires-Dist: monai>=1.3.2; extra == "vision"
 Requires-Dist: einops>=0.8.1; extra == "vision"
+Provides-Extra: language
+Requires-Dist: datasets>=3.2.0; extra == "language"
+Requires-Dist: litellm>=1.61.8; extra == "language"
 Provides-Extra: all
 Requires-Dist: h5py>=3.10.0; extra == "all"
 Requires-Dist: nibabel>=4.0.1; extra == "all"
@@ -255,6 +258,8 @@ Requires-Dist: imagesize>=1.4.1; extra == "all"
 Requires-Dist: scipy>=1.14.0; extra == "all"
 Requires-Dist: monai>=1.3.2; extra == "all"
 Requires-Dist: einops>=0.8.1; extra == "all"
+Requires-Dist: datasets>=3.2.0; extra == "all"
+Requires-Dist: litellm>=1.61.8; extra == "all"
 Description-Content-Type: text/markdown
 
 <div align="center">
@@ -293,7 +298,7 @@ Check out the [documentation](https://kaiko-ai.github.io/eva/) for more informat
 
 ### Highlights:
 - Easy and reliable benchmark of Oncology FMs
-- Supports path-level classification, slide-level classification and semantic segmentation downstream tasks
+- Supports patch-level classification, slide-level classification, semantic segmentation, and text classification downstream tasks
 - Automatic embedding inference and evaluation of a downstream task
 - Native support of popular medical [datasets](https://kaiko-ai.github.io/eva/dev/datasets/) and models
 - Produce statistics over multiple evaluation fits and multiple metrics
@@ -308,6 +313,9 @@ pip install kaiko-eva
 # to install the expanded `vision` version
 pip install 'kaiko-eva[vision]'
 
+# to install the expanded `language` version
+pip install 'kaiko-eva[language]'
+
 # to install everything
 pip install 'kaiko-eva[all]'
 ```
@@ -448,26 +456,28 @@ input, resulting in a faster evaluation.
 Here are some examples to get you started:
 
 - Perform a downstream offline **classification** evaluation of `DINO ViT-S/16`
-on the `BACH` dataset with linear probing by first inferring the embeddings
-and then performing 5 sequential fits:
+on the `BACH` dataset with linear probing by first pre-calculating the embeddings:
   ```sh
-  export DOWNLOAD_DATA=true
-  eva predict_fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/dino_vit/offline/bach.yaml
+  DOWNLOAD_DATA=true \
+  MODEL_NAME=universal/vit_small_patch16_224_dino \
+  eva predict_fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/pathology/offline/classification/bach.yaml
   ```
 
-- Perform a downstream online **segmentation** evaluation of `DINO ViT-S/16` on the
-`MoNuSAC` dataset with the `ConvDecoderMS` decoder:
+- Perform a downstream online **segmentation** evaluation of `DINO ViT-S/16` on the `MoNuSAC` dataset with the `ConvDecoderWithImage` decoder:
   ```sh
-  export DOWNLOAD_DATA=true
-  eva fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/dino_vit/online/monusac.yaml
+  DOWNLOAD_DATA=true \
+  MODEL_NAME=universal/vit_small_patch16_224_dino \
+  eva fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/pathology/online/segmentation/monusac.yaml
   ```
 
+By default `eva` will perform 5 evaluation runs using different seeds, however, you can control the number of runs through the `N_RUNS` environment variable or in the configuration file. The results will be saved to `./logs` by default, or to `OUTPUT_ROOT` if specified.
+
 For more examples, take a look at the [configs](https://github.com/kaiko-ai/eva/tree/main/configs)
-and [tutorials](https://kaiko-ai.github.io/eva/dev/user-guide/advanced/replicate_evaluations/).
+and [tutorials](https://kaiko-ai.github.io/eva/main/user-guide/advanced/replicate_evaluations/).
 
 > [!NOTE]
 > All the datasets that support automatic download in the repo have by default the option to automatically download set to false.
-> For automatic download you have to manually set the environmental variable `DOWNLOAD_DATA=true` or in the configuration file `download=true`.
+> For automatic download you have to manually set the environment variable `DOWNLOAD_DATA=true` or in the configuration file `download=true`.
 
 ## Leaderboards
 

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/README.md

@@ -34,7 +34,7 @@ Check out the [documentation](https://kaiko-ai.github.io/eva/) for more informat
 
 ### Highlights:
 - Easy and reliable benchmark of Oncology FMs
-- Supports path-level classification, slide-level classification and semantic segmentation downstream tasks
+- Supports patch-level classification, slide-level classification, semantic segmentation, and text classification downstream tasks
 - Automatic embedding inference and evaluation of a downstream task
 - Native support of popular medical [datasets](https://kaiko-ai.github.io/eva/dev/datasets/) and models
 - Produce statistics over multiple evaluation fits and multiple metrics
@@ -49,6 +49,9 @@ pip install kaiko-eva
 # to install the expanded `vision` version
 pip install 'kaiko-eva[vision]'
 
+# to install the expanded `language` version
+pip install 'kaiko-eva[language]'
+
 # to install everything
 pip install 'kaiko-eva[all]'
 ```
@@ -189,26 +192,28 @@ input, resulting in a faster evaluation.
 Here are some examples to get you started:
 
 - Perform a downstream offline **classification** evaluation of `DINO ViT-S/16`
-on the `BACH` dataset with linear probing by first inferring the embeddings
-and then performing 5 sequential fits:
+on the `BACH` dataset with linear probing by first pre-calculating the embeddings:
   ```sh
-  export DOWNLOAD_DATA=true
-  eva predict_fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/dino_vit/offline/bach.yaml
+  DOWNLOAD_DATA=true \
+  MODEL_NAME=universal/vit_small_patch16_224_dino \
+  eva predict_fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/pathology/offline/classification/bach.yaml
   ```
 
-- Perform a downstream online **segmentation** evaluation of `DINO ViT-S/16` on the
-`MoNuSAC` dataset with the `ConvDecoderMS` decoder:
+- Perform a downstream online **segmentation** evaluation of `DINO ViT-S/16` on the `MoNuSAC` dataset with the `ConvDecoderWithImage` decoder:
   ```sh
-  export DOWNLOAD_DATA=true
-  eva fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/dino_vit/online/monusac.yaml
+  DOWNLOAD_DATA=true \
+  MODEL_NAME=universal/vit_small_patch16_224_dino \
+  eva fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/pathology/online/segmentation/monusac.yaml
   ```
 
+By default `eva` will perform 5 evaluation runs using different seeds, however, you can control the number of runs through the `N_RUNS` environment variable or in the configuration file. The results will be saved to `./logs` by default, or to `OUTPUT_ROOT` if specified.
+
 For more examples, take a look at the [configs](https://github.com/kaiko-ai/eva/tree/main/configs)
-and [tutorials](https://kaiko-ai.github.io/eva/dev/user-guide/advanced/replicate_evaluations/).
+and [tutorials](https://kaiko-ai.github.io/eva/main/user-guide/advanced/replicate_evaluations/).
 
 > [!NOTE]
 > All the datasets that support automatic download in the repo have by default the option to automatically download set to false.
-> For automatic download you have to manually set the environmental variable `DOWNLOAD_DATA=true` or in the configuration file `download=true`.
+> For automatic download you have to manually set the environment variable `DOWNLOAD_DATA=true` or in the configuration file `download=true`.
 
 ## Leaderboards
 

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "kaiko-eva"
-version = "0.2.1"
+version = "0.3.0"
 description = "Evaluation Framework for oncology foundation models."
 keywords = [
     "machine-learning",
@@ -72,6 +72,10 @@ vision = [
     "monai>=1.3.2",
     "einops>=0.8.1",
 ]
+language = [
+    "datasets>=3.2.0",
+    "litellm>=1.61.8",
+]
 all = [
     "h5py>=3.10.0",
     "nibabel>=4.0.1",
@@ -85,6 +89,8 @@ all = [
     "scipy>=1.14.0",
     "monai>=1.3.2",
     "einops>=0.8.1",
+    "datasets>=3.2.0",
+    "litellm>=1.61.8",
 ]
 
 [project.scripts]

kaiko_eva-0.3.0/src/eva/core/data/dataloaders/__init__.py

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

kaiko_eva-0.3.0/src/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"]

kaiko_eva-0.3.0/src/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)

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/src/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,
         )

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/src/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,
+        )

kaiko_eva-0.3.0/src/eva/core/metrics/defaults/classification/multiclass.py

@@ -0,0 +1,82 @@
+"""Default metric collection for multiclass classification tasks."""
+
+from typing import Literal
+
+from torchmetrics import classification
+
+from eva.core.metrics import structs
+
+
+class MulticlassClassificationMetrics(structs.MetricCollection):
+    """Default metrics for multi-class classification tasks."""
+
+    def __init__(
+        self,
+        num_classes: int,
+        average: Literal["macro", "weighted", "none"] = "macro",
+        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.
+
+        Args:
+            num_classes: Integer specifying the number of classes.
+            average: Defines the reduction that is applied over labels.
+            ignore_index: Specifies a target value that is ignored and does not
+                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.
+        """
+        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,
+                )
+            )
+            compute_groups.append(["MulticlassAUROC"])
+
+        super().__init__(
+            metrics=metrics,
+            prefix=prefix,
+            postfix=postfix,
+            compute_groups=compute_groups,
+        )

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/src/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"]

kaiko_eva-0.3.0/src/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,
+        }

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/src/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:

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/src/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)
 

kaiko_eva-0.3.0/src/eva/core/models/wrappers/base.py

@@ -0,0 +1,51 @@
+"""Base class for model wrappers."""
+
+import abc
+from typing import Callable, Generic, TypeVar
+
+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, Generic[InputType, OutputType]):
+    """Base class for model wrappers."""
+
+    def __init__(self, transforms: Callable | None = None) -> None:
+        """Initializes the model.
+
+        Args:
+            transforms: The transforms to apply to the output produced by the model.
+        """
+        super().__init__()
+
+        self._output_transforms = transforms
+
+        self._model: Callable[..., OutputType] | nn.Module
+
+    @override
+    def forward(self, tensor: InputType) -> OutputType:
+        out = self.model_forward(tensor)
+        return self._apply_transforms(out)
+
+    @abc.abstractmethod
+    def load_model(self) -> Callable[..., OutputType]:
+        """Loads the model."""
+        raise NotImplementedError
+
+    def model_forward(self, tensor: InputType) -> OutputType:
+        """Implements the forward pass of the model.
+
+        Args:
+            tensor: The input tensor to the model.
+        """
+        return self._model(tensor)
+
+    def _apply_transforms(self, tensor: OutputType) -> OutputType:
+        if self._output_transforms is not None:
+            tensor = self._output_transforms(tensor)
+        return tensor

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/src/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

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/src/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

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/src/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
         )

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/src/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

{kaiko_eva-0.2.1 → kaiko_eva-0.3.0}/src/eva/core/trainers/_recorder.py

@@ -129,7 +129,10 @@ class SessionRecorder:
     def _save_config(self) -> None:
         """Saves the config yaml with resolved env placeholders to the output directory."""
         if self.config_path:
-            config = OmegaConf.load(self.config_path)
+            config_fs = cloud_io.get_filesystem(self.config_path)
+            with config_fs.open(self.config_path, "r") as config_file:
+                config = OmegaConf.load(config_file)  # type: ignore
+
             fs = cloud_io.get_filesystem(self._output_dir, anon=False)
             with fs.open(os.path.join(self._output_dir, self._config_file), "w") as file:
                 config_yaml = OmegaConf.to_yaml(config, resolve=True)