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

eva/core/callbacks/config.py

@@ -51,6 +51,10 @@ class ConfigurationLogger(pl.Callback):
 
         save_as = os.path.join(log_dir, self._save_as)
         fs = cloud_io.get_filesystem(log_dir)
+
+        if not fs.exists(log_dir):
+            fs.makedirs(log_dir)
+
         with fs.open(save_as, "w") as output_file:
             yaml.dump(configuration, output_file, sort_keys=False)
 

eva/core/cli/setup.py

@@ -59,7 +59,7 @@ def _initialize_logger() -> None:
         " :: <bold><level>{level}</level></bold>"
         " :: {message}",
         colorize=True,
-        level="INFO",
+        level=os.getenv("LOGURU_LEVEL", "INFO"),
     )
 
 

eva/core/data/dataloaders/__init__.py

@@ -1,6 +1,5 @@
 """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"]
+__all__ = ["DataLoader"]

eva/core/data/samplers/random.py

@@ -2,6 +2,7 @@
 
 from typing import Optional
 
+import torch
 from torch.utils import data
 from typing_extensions import override
 
@@ -10,30 +11,36 @@ from eva.core.data.samplers.sampler import SamplerWithDataSource
 
 
 class RandomSampler(data.RandomSampler, SamplerWithDataSource[int]):
-    """Samples elements randomly."""
+    """Samples elements randomly from a MapDataset."""
 
     data_source: datasets.MapDataset  # type: ignore
 
     def __init__(
-        self, replacement: bool = False, num_samples: Optional[int] = None, generator=None
+        self,
+        replacement: bool = False,
+        num_samples: Optional[int] = None,
+        seed: Optional[int] = None,
     ) -> None:
-        """Initializes the random sampler.
+        """Initialize the random sampler.
 
         Args:
-            data_source: dataset to sample from
-            replacement: samples are drawn on-demand with replacement if ``True``, default=``False``
-            num_samples: number of samples to draw, default=`len(dataset)`.
-            generator: Generator used in sampling.
+            replacement: Samples are drawn on-demand with replacement if ``True``, default=``False``
+            num_samples: Number of samples to draw, default=``len(dataset)``.
+            seed: Optional seed for the random number generator.
         """
         self.replacement = replacement
         self._num_samples = num_samples
-        self.generator = generator
+        self._generator = None
+
+        if seed is not None:
+            self._generator = torch.Generator()
+            self._generator.manual_seed(seed)
 
     @override
     def set_dataset(self, data_source: datasets.MapDataset) -> None:
         super().__init__(
             data_source,
             replacement=self.replacement,
-            num_samples=self.num_samples,
-            generator=self.generator,
+            num_samples=self._num_samples,
+            generator=self._generator,
         )

eva/core/interface/interface.py

@@ -132,3 +132,24 @@ class Interface:
             n_runs=trainer.n_runs,
             verbose=trainer.n_runs > 1,
         )
+
+    def validate_test(
+        self,
+        trainer: eva_trainer.Trainer,
+        model: modules.ModelModule,
+        data: datamodules.DataModule,
+    ) -> None:
+        """Runs validation & test stages."""
+        if getattr(data.datasets, "val", None) is None:
+            raise ValueError("The provided data module does not contain a validation dataset.")
+        if getattr(data.datasets, "test", None) is None:
+            raise ValueError("The provided data module does not contain a test dataset.")
+
+        eva_trainer.run_evaluation_session(
+            base_trainer=trainer,
+            base_model=model,
+            datamodule=data,
+            stages=["validate", "test"],
+            n_runs=trainer.n_runs,
+            verbose=trainer.n_runs > 1,
+        )

eva/core/models/modules/module.py

@@ -33,8 +33,8 @@ class ModelModule(pl.LightningModule):
         super().__init__()
 
         self._metrics = metrics or self.default_metrics
-        self._postprocess = postprocess or self.default_postprocess
 
+        self.postprocess = postprocess or self.default_postprocess
         self.metrics = metrics_lib.MetricModule.from_schema(self._metrics)
 
     @property
@@ -133,7 +133,7 @@ class ModelModule(pl.LightningModule):
         Returns:
             The updated outputs.
         """
-        self._postprocess(outputs)
+        self.postprocess(outputs)
         return memory.recursive_detach(outputs, to_cpu=self.metrics_device.type == "cpu")
 
     def _forward_and_log_metrics(

eva/core/models/wrappers/base.py

@@ -25,7 +25,7 @@ class BaseModel(nn.Module, Generic[InputType, OutputType]):
 
         self._output_transforms = transforms
 
-        self._model: Callable[..., OutputType] | nn.Module
+        self.model: Callable[..., OutputType] | nn.Module
 
     @override
     def forward(self, tensor: InputType) -> OutputType:
@@ -43,7 +43,7 @@ class BaseModel(nn.Module, Generic[InputType, OutputType]):
         Args:
             tensor: The input tensor to the model.
         """
-        return self._model(tensor)
+        return self.model(tensor)
 
     def _apply_transforms(self, tensor: OutputType) -> OutputType:
         if self._output_transforms is not None:

eva/core/models/wrappers/from_function.py

@@ -41,12 +41,12 @@ class ModelFromFunction(base.BaseModel[torch.Tensor, torch.Tensor]):
         self._arguments = arguments
         self._checkpoint_path = checkpoint_path
 
-        self.load_model()
+        self.model = self.load_model()
 
     @override
-    def load_model(self) -> None:
+    def load_model(self) -> nn.Module:
         class_path = jsonargparse.class_from_function(self._path, func_return=nn.Module)
         model = class_path(**self._arguments or {})
         if self._checkpoint_path is not None:
             _utils.load_model_weights(model, self._checkpoint_path)
-        self._model = model
+        return model

eva/core/models/wrappers/from_torchhub.py

@@ -52,12 +52,12 @@ class TorchHubModel(base.BaseModel[torch.Tensor, torch.Tensor]):
         self._trust_repo = trust_repo
         self._model_kwargs = model_kwargs or {}
 
-        self.load_model()
+        self.model = self.load_model()
 
     @override
-    def load_model(self) -> None:
+    def load_model(self) -> nn.Module:
         """Builds and loads the torch.hub model."""
-        self._model: nn.Module = torch.hub.load(
+        model: nn.Module = torch.hub.load(
             repo_or_dir=self._repo_or_dir,
             model=self._model_name,
             trust_repo=self._trust_repo,
@@ -66,21 +66,23 @@ class TorchHubModel(base.BaseModel[torch.Tensor, torch.Tensor]):
         )  # type: ignore
 
         if self._checkpoint_path:
-            _utils.load_model_weights(self._model, self._checkpoint_path)
+            _utils.load_model_weights(model, self._checkpoint_path)
 
         TorchHubModel.__name__ = self._model_name
 
+        return model
+
     @override
     def model_forward(self, tensor: torch.Tensor) -> torch.Tensor | List[torch.Tensor]:
         if self._out_indices is not None:
-            if not hasattr(self._model, "get_intermediate_layers"):
+            if not hasattr(self.model, "get_intermediate_layers"):
                 raise ValueError(
                     "Only models with `get_intermediate_layers` are supported "
                     "when using `out_indices`."
                 )
 
             return list(
-                self._model.get_intermediate_layers(
+                self.model.get_intermediate_layers(  # type: ignore
                     tensor,
                     self._out_indices,
                     reshape=True,
@@ -89,4 +91,4 @@ class TorchHubModel(base.BaseModel[torch.Tensor, torch.Tensor]):
                 )
             )
 
-        return self._model(tensor)
+        return self.model(tensor)

eva/core/models/wrappers/huggingface.py

@@ -4,6 +4,7 @@ from typing import Any, Callable, Dict
 
 import torch
 import transformers
+from torch import nn
 from typing_extensions import override
 
 from eva.core.models.wrappers import base
@@ -33,12 +34,10 @@ class HuggingFaceModel(base.BaseModel[torch.Tensor, torch.Tensor]):
         self._model_name_or_path = model_name_or_path
         self._model_kwargs = model_kwargs or {}
 
-        self.load_model()
+        self.model = self.load_model()
 
     @override
-    def load_model(self) -> None:
+    def load_model(self) -> nn.Module:
         # 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, **model_kwargs
-        )
+        return transformers.AutoModel.from_pretrained(self._model_name_or_path, **model_kwargs)

eva/core/models/wrappers/onnx.py

@@ -30,21 +30,21 @@ class ONNXModel(base.BaseModel[torch.Tensor, torch.Tensor]):
         self._path = path
         self._device = device
 
-        self.load_model()
+        self.model = self.load_model()
 
     @override
     def load_model(self) -> Any:
         if self._device == "cuda" and not torch.cuda.is_available():
             raise ValueError("Device is set to 'cuda', but CUDA is not available.")
         provider = "CUDAExecutionProvider" if self._device == "cuda" else "CPUExecutionProvider"
-        self._model = ort.InferenceSession(self._path, providers=[provider])  # type: ignore
+        return ort.InferenceSession(self._path, providers=[provider])  # type: ignore
 
     @override
     def model_forward(self, tensor: torch.Tensor) -> torch.Tensor:
         # TODO: Use IO binding to avoid copying the tensor to CPU.
         # https://onnxruntime.ai/docs/api/python/api_summary.html#data-on-device
-        if not isinstance(self._model, ort.InferenceSession):
+        if not isinstance(self.model, ort.InferenceSession):
             raise ValueError("Model is not loaded.")
-        inputs = {self._model.get_inputs()[0].name: tensor.detach().cpu().numpy()}
-        outputs = self._model.run(None, inputs)[0]
+        inputs = {self.model.get_inputs()[0].name: tensor.detach().cpu().numpy()}
+        outputs = self.model.run(None, inputs)[0]
         return torch.from_numpy(outputs).float().to(tensor.device)

eva/core/trainers/trainer.py

@@ -8,6 +8,7 @@ from lightning.pytorch import loggers as pl_loggers
 from lightning.pytorch import trainer as pl_trainer
 from lightning.pytorch.utilities import argparse
 from lightning_fabric.utilities import cloud_io
+from lightning_utilities.core.rank_zero import rank_zero_only
 from typing_extensions import override
 
 from eva.core import loggers as eva_loggers
@@ -66,6 +67,7 @@ class Trainer(pl_trainer.Trainer):
     def log_dir(self) -> str | None:
         return self.strategy.broadcast(self._log_dir)
 
+    @rank_zero_only
     def init_logger_run(self, run_id: int | None) -> None:
         """Setup the loggers & log directories when starting a new run.
 

eva/language/__init__.py

@@ -1,6 +1,7 @@
 """eva language API."""
 
 try:
+    from eva.language import models
     from eva.language.data import datasets
 except ImportError as e:
     msg = (
@@ -10,4 +11,4 @@ except ImportError as e:
     )
     raise ImportError(str(e) + "\n\n" + msg) from e
 
-__all__ = ["datasets"]
+__all__ = ["models", "datasets"]

eva/language/callbacks/__init__.py

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

eva/language/callbacks/writers/__init__.py

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

eva/language/callbacks/writers/prediction.py

@@ -0,0 +1,176 @@
+"""Text prediction writer callbacks."""
+
+import abc
+import os
+from typing import Any, Dict, List, Literal, Sequence, Tuple, TypedDict
+
+import lightning.pytorch as pl
+import pandas as pd
+import torch
+from lightning.pytorch import callbacks
+from torch import nn
+from typing_extensions import NotRequired, override
+
+from eva.core.models.modules import utils as module_utils
+from eva.language.models.typings import TextBatch
+from eva.language.utils.text import messages as message_utils
+
+
+class ManifestEntry(TypedDict):
+    """A single entry in the manifest file."""
+
+    prediction: str
+    """The predicted text."""
+
+    target: str
+    """The ground truth text."""
+
+    text: NotRequired[str]
+    """The input text data."""
+
+    split: NotRequired[str]
+    """The dataset split (e.g. train, val, test)."""
+
+
+class TextPredictionWriter(callbacks.BasePredictionWriter, abc.ABC):
+    """Callback for writing generated text predictions to disk."""
+
+    def __init__(
+        self,
+        output_dir: str,
+        model: nn.Module,
+        dataloader_idx_map: Dict[int, str] | None = None,
+        metadata_keys: List[str] | None = None,
+        include_input: bool = True,
+        overwrite: bool = False,
+        apply_postprocess: bool = False,
+        save_format: Literal["jsonl", "parquet", "csv"] = "jsonl",
+    ) -> None:
+        """Initializes a new callback.
+
+        Args:
+            output_dir: The directory where the embeddings will be saved.
+            model: The model instance used to generate the predictions.
+            dataloader_idx_map: A dictionary mapping dataloader indices to their respective
+                names (e.g. train, val, test).
+            metadata_keys: An optional list of keys to extract from the batch metadata and store
+                as additional columns in the manifest file.
+            include_input: Whether to include the original input text messages in the output.
+            overwrite: Whether to overwrite if embeddings are already present in the specified
+                output directory. If set to `False`, an error will be raised if embeddings are
+                already present (recommended).
+            apply_postprocess: Whether to apply the postprocesses specified in the model module.
+            save_format: The file format to use for saving the manifest file with the predictions.
+        """
+        super().__init__()
+        self.output_dir = output_dir
+        self.model = model
+        self.dataloader_idx_map = dataloader_idx_map or {}
+        self.metadata_keys = metadata_keys
+        self.include_input = include_input
+        self.overwrite = overwrite
+        self.apply_postprocess = apply_postprocess
+        self.save_format = save_format
+
+        self._manifest_path = os.path.join(self.output_dir, f"manifest.{self.save_format}")
+        self._data: List[ManifestEntry] = []
+
+    @override
+    def on_predict_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        self._check_if_exists()
+
+        self.model = self.model.to(pl_module.device)
+        self.model.eval()
+
+    @override
+    def write_on_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        prediction: Any,
+        batch_indices: Sequence[int],
+        batch: TextBatch,
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        text_batch, target_batch, metadata_batch = self._unpack_batch(batch)
+        has_target = target_batch is not None
+        split = self.dataloader_idx_map.get(dataloader_idx, "")
+
+        prediction_batch = self._get_predictions(batch)
+
+        target_batch, prediction_batch = self._apply_postprocess(
+            pl_module, target_batch, prediction_batch
+        )
+
+        for i in range(len(batch_indices)):
+            entry: ManifestEntry = {
+                "text": message_utils.serialize(text_batch[i]),
+                "prediction": str(prediction_batch[i]),
+                "target": str(target_batch[i]) if has_target else "",
+                "split": split if split else "",
+            }
+
+            if self.metadata_keys is not None and metadata_batch is not None:
+                for key in self.metadata_keys:
+                    entry[key] = metadata_batch[key][i]
+
+            self._data.append(entry)
+
+    @override
+    def on_predict_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        """Saves the gathered predictions to a manifest file."""
+        df = pd.DataFrame(self._data)
+
+        match self.save_format:
+            case "jsonl":
+                df.to_json(self._manifest_path, orient="records", lines=True)
+            case "parquet":
+                df.to_parquet(self._manifest_path, index=False)
+            case "csv":
+                df.to_csv(self._manifest_path, index=False)
+            case _:
+                raise ValueError(f"Unsupported save format: {self.save_format}")
+
+    def _get_predictions(self, batch: TextBatch) -> List[str]:
+        with torch.no_grad():
+            predictions = self.model(batch)
+
+        if not isinstance(predictions, list) or not all(isinstance(p, str) for p in predictions):
+            raise ValueError("The model's output should be a list of strings.")
+
+        return predictions
+
+    def _check_if_exists(self) -> None:
+        """Checks if the output directory already exists and if it should be overwritten."""
+        os.makedirs(self.output_dir, exist_ok=True)
+        if os.path.exists(self._manifest_path) and not self.overwrite:
+            raise FileExistsError(
+                f"The specified output directory already exists: {self.output_dir}. This "
+                "either means that the predictions have been computed before or that a "
+                "wrong output directory is being used."
+            )
+        os.makedirs(self.output_dir, exist_ok=True)
+
+    def _apply_postprocess(
+        self, pl_module: pl.LightningModule, targets: Any, predictions: Any
+    ) -> Tuple[List[Any], List[Any]]:
+        def _to_list(data: Any) -> List[Any]:
+            if isinstance(data, torch.Tensor):
+                return data.cpu().tolist()
+            return data
+
+        if self.apply_postprocess and hasattr(pl_module, "postprocess"):
+            if (
+                isinstance(pl_module.postprocess, module_utils.BatchPostProcess)
+                and pl_module.postprocess.predictions_transforms is not None
+            ):
+                outputs = {"targets": targets, "predictions": predictions}
+                pl_module.postprocess(outputs)
+                targets, predictions = outputs["targets"], outputs["predictions"]
+
+        return _to_list(targets), _to_list(predictions)
+
+    def _unpack_batch(self, batch: TextBatch) -> Tuple[list, list | None, dict | None]:
+        text_batch, target_batch, metadata_batch = TextBatch(*batch)
+        return text_batch, target_batch, metadata_batch

eva/language/data/dataloaders/__init__.py

@@ -0,0 +1,5 @@
+"""Language Dataloaders API."""
+
+from eva.language.data.dataloaders.collate_fn import prediction_collate, text_collate
+
+__all__ = ["text_collate", "prediction_collate"]

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

@@ -0,0 +1,5 @@
+"""Collate functions API."""
+
+from eva.language.data.dataloaders.collate_fn.text import prediction_collate, text_collate
+
+__all__ = ["text_collate", "prediction_collate"]

eva/language/data/dataloaders/collate_fn/text.py

@@ -0,0 +1,57 @@
+"""Collate functions for text data."""
+
+from typing import List
+
+from torch.utils.data._utils.collate import default_collate
+
+from eva.language.data.datasets.typings import PredictionSample, TextSample
+from eva.language.models.typings import PredictionBatch, TextBatch
+
+
+def text_collate(batch: List[TextSample]) -> TextBatch:
+    """Collate function for text data that keeps texts as separate strings.
+
+    Args:
+        batch: List of tuples containing (text, target, metadata) from the dataset
+
+    Returns:
+        A batch of text samples with targets and metadata.
+    """
+    texts, targets, metadata = zip(*batch, strict=False)
+    first_sample = batch[0]
+    metadata = None
+    if first_sample.metadata is not None:
+        metadata = {
+            k: [sample.metadata[k] for sample in batch if sample.metadata]
+            for k in first_sample.metadata.keys()
+        }
+    return TextBatch(
+        text=list(texts),
+        target=default_collate(targets) if targets[0] is not None else None,
+        metadata=metadata,
+    )
+
+
+def prediction_collate(batch: List[PredictionSample]) -> PredictionBatch:
+    """Collate function for text prediction data.
+
+    Args:
+        batch: List of tuples containing (prediction, target, text, metadata) from the dataset
+
+    Returns:
+        A batch of prediction samples.
+    """
+    predictions, targets, texts, metadata = zip(*batch, strict=False)
+    first_sample = batch[0]
+    metadata = None
+    if first_sample.metadata is not None:
+        metadata = {
+            k: [sample.metadata[k] for sample in batch if sample.metadata]
+            for k in first_sample.metadata.keys()
+        }
+    return PredictionBatch(
+        prediction=default_collate(predictions) if predictions[0] is not None else None,
+        target=default_collate(targets) if targets[0] is not None else None,
+        text=list(texts) if first_sample.text is not None else None,
+        metadata=metadata,
+    )

eva/language/data/datasets/__init__.py

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

eva/language/data/datasets/{language.py → base.py}

@@ -10,4 +10,4 @@ DataSample = TypeVar("DataSample")
 
 
 class LanguageDataset(base.MapDataset, abc.ABC, Generic[DataSample]):
-    """Base dataset class for text tasks."""
+    """Base dataset class for language tasks."""

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

@@ -1,15 +1,13 @@
 """Base for text classification datasets."""
 
-import abc
-from typing import Any, Dict, List, Tuple
+from typing import Dict, List
 
 import torch
-from typing_extensions import override
 
-from eva.language.data.datasets.language import LanguageDataset
+from eva.language.data.datasets.text import TextDataset
 
 
-class TextClassification(LanguageDataset[Tuple[str, torch.Tensor, Dict[str, Any]]], abc.ABC):
+class TextClassification(TextDataset[torch.Tensor]):
     """Text classification abstract dataset."""
 
     def __init__(self) -> None:
@@ -23,41 +21,3 @@ class TextClassification(LanguageDataset[Tuple[str, torch.Tensor, Dict[str, Any]
     @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 {})