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

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,201 @@
+"""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
+import torch.distributed as dist
+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.core.utils import distributed as dist_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] = []
+        self._is_rank_zero: bool = False
+
+    @override
+    def on_predict_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        self._is_rank_zero = trainer.is_global_zero
+
+        if self._is_rank_zero:
+            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 = {
+                "prediction": str(prediction_batch[i]),
+                "target": str(target_batch[i]) if has_target else "",
+                "split": split if split else "",
+            }
+            if self.include_input:
+                entry["text"] = message_utils.serialize(text_batch[i])
+
+            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."""
+        if dist_utils.is_distributed():
+            dist.barrier()
+            data = self._gather_data_from_ranks()
+        else:
+            data = self._data
+
+        if self._is_rank_zero:
+            df = pd.DataFrame(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 _gather_data_from_ranks(self) -> List[ManifestEntry]:
+        world_size = dist.get_world_size()
+        gathered: List[List[ManifestEntry] | None] = [None] * world_size
+        dist.all_gather_object(gathered, self._data)
+        return [row for shard in gathered for row in (shard or [])]
+
+    def _get_predictions(self, batch: TextBatch) -> List[str]:
+        with torch.no_grad():
+            output = self.model(batch)
+
+        if (
+            not isinstance(output, dict)
+            or "generated_text" not in output
+            or not all(isinstance(p, str) for p in output["generated_text"])
+        ):
+            raise ValueError(
+                f"A dictionary with 'generated_text' key is expected, got {type(output)}"
+            )
+
+        return output["generated_text"]
+
+    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."
+            )
+
+    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 {})

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

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

eva/language/data/datasets/prediction.py

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

eva/language/data/datasets/schemas.py

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