kaiko-eva 0.4.0__tar.gz → 0.4.1__tar.gz

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kaiko-eva
-Version: 0.4.0
+Version: 0.4.1
 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>
@@ -491,7 +491,13 @@ and [tutorials](https://kaiko-ai.github.io/eva/main/user-guide/advanced/replicat
 
 The following table shows the FMs we have evaluated with _`eva`_. For more detailed information about the evaluation process, please refer to our [documentation](https://kaiko-ai.github.io/eva/main/leaderboards/).
 
-![Pathology Leaderboard](./docs/images/leaderboard.svg)
+### Pathology
+
+<img src="./docs/images/leaderboards/pathology.svg" alt="Pathology Leaderboard">
+
+### Radiology
+
+<img src="./docs/images/leaderboards/radiology.svg" alt="Radiology Leaderboard" width="600">
 
 
 ## Contributing

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/README.md

@@ -222,7 +222,13 @@ and [tutorials](https://kaiko-ai.github.io/eva/main/user-guide/advanced/replicat
 
 The following table shows the FMs we have evaluated with _`eva`_. For more detailed information about the evaluation process, please refer to our [documentation](https://kaiko-ai.github.io/eva/main/leaderboards/).
 
-![Pathology Leaderboard](./docs/images/leaderboard.svg)
+### Pathology
+
+<img src="./docs/images/leaderboards/pathology.svg" alt="Pathology Leaderboard">
+
+### Radiology
+
+<img src="./docs/images/leaderboards/radiology.svg" alt="Radiology Leaderboard" width="600">
 
 
 ## Contributing

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "kaiko-eva"
-version = "0.4.0"
+version = "0.4.1"
 description = "Evaluation Framework for oncology foundation models."
 keywords = [
     "machine-learning",

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/core/callbacks/config.py

@@ -9,11 +9,13 @@ from typing import Any, Dict, List
 import lightning.pytorch as pl
 import yaml
 from lightning_fabric.utilities import cloud_io
+from loguru import logger
 from loguru import logger as cli_logger
 from omegaconf import OmegaConf
 from typing_extensions import TypeGuard, override
 
 from eva.core import loggers
+from eva.core.utils import distributed as dist_utils
 
 
 class ConfigurationLogger(pl.Callback):
@@ -39,8 +41,14 @@ class ConfigurationLogger(pl.Callback):
         pl_module: pl.LightningModule,
         stage: str | None = None,
     ) -> None:
-        log_dir = trainer.log_dir
-        if not _logdir_exists(log_dir):
+        if dist_utils.is_distributed():
+            logger.info("ConfigurationLogger skipped as not supported in distributed mode.")
+            # TODO: Enabling leads to deadlocks in DDP mode, but I could not yet figure out why.
+            return
+
+        if not trainer.is_global_zero or not _logdir_exists(
+            log_dir := trainer.log_dir, self._verbose
+        ):
             return
 
         configuration = _load_submitted_config()
@@ -130,7 +138,7 @@ def _type_resolver(mapping: Dict[str, Any]) -> Dict[str, Any]:
     for key, value in mapping.items():
         if isinstance(value, dict):
             formatted_value = _type_resolver(value)
-        elif isinstance(value, list) and isinstance(value[0], dict):
+        elif isinstance(value, list) and value and isinstance(value[0], dict):
             formatted_value = [_type_resolver(subvalue) for subvalue in value]
         else:
             try:
@@ -138,10 +146,7 @@ def _type_resolver(mapping: Dict[str, Any]) -> Dict[str, Any]:
                 formatted_value = (
                     value if isinstance(parsed_value, BuiltinFunctionType) else parsed_value
                 )
-
             except Exception:
                 formatted_value = value
-
         mapping[key] = formatted_value
-
     return mapping

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/core/callbacks/writers/embeddings/base.py

@@ -7,6 +7,7 @@ from typing import Any, Dict, List, Sequence
 
 import lightning.pytorch as pl
 import torch
+import torch.distributed as dist
 from lightning.pytorch import callbacks
 from loguru import logger
 from torch import multiprocessing, nn
@@ -15,6 +16,7 @@ from typing_extensions import override
 from eva.core import utils
 from eva.core.callbacks.writers.embeddings.typings import QUEUE_ITEM
 from eva.core.models.modules.typings import INPUT_BATCH
+from eva.core.utils import distributed as dist_utils
 from eva.core.utils import multiprocessing as eva_multiprocessing
 
 
@@ -58,8 +60,9 @@ class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
         self._save_every_n = save_every_n
         self._metadata_keys = metadata_keys or []
 
-        self._write_queue: multiprocessing.Queue
-        self._write_process: eva_multiprocessing.Process
+        self._write_queue: multiprocessing.Queue | None = None
+        self._write_process: eva_multiprocessing.Process | None = None
+        self._is_rank_zero: bool = False
 
     @staticmethod
     @abc.abstractmethod
@@ -78,9 +81,13 @@ class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
 
     @override
     def on_predict_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
-        self._check_if_exists()
-        self._initialize_write_process()
-        self._write_process.start()
+        self._is_rank_zero = trainer.is_global_zero
+        if self._is_rank_zero:
+            self._check_if_exists()
+            self._initialize_write_process()
+            if self._write_process is None or self._write_queue is None:
+                raise RuntimeError("Failed to initialize embedding writer process.")
+            self._write_process.start()
 
         if self._backbone is not None:
             self._backbone = self._backbone.to(pl_module.device)
@@ -106,6 +113,7 @@ class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
         with torch.no_grad():
             embeddings = self._get_embeddings(prediction)
 
+        queue_items: List[QUEUE_ITEM] = []
         for local_idx, global_idx in enumerate(batch_indices[: len(embeddings)]):
             data_name = dataset.filename(global_idx)
             save_name = os.path.splitext(data_name)[0] + ".pt"
@@ -121,15 +129,41 @@ class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
                 split=split,
                 metadata=item_metadata,
             )
-            self._write_queue.put(item)
+            queue_items.append(item)
 
-        self._write_process.check_exceptions()
+        gathered_items = self._gather_queue_items(queue_items)
+        if self._is_rank_zero:
+            for item in gathered_items:
+                self._write_queue.put(item)  # type: ignore
+                self._write_process.check_exceptions()  # type: ignore
 
     @override
     def on_predict_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
-        self._write_queue.put(None)
-        self._write_process.join()
-        logger.info(f"Predictions and manifest saved to {self._output_dir}")
+        if dist_utils.is_distributed():
+            dist.barrier()
+
+        if self._is_rank_zero and self._write_queue is not None:
+            self._write_queue.put(None)
+            if self._write_process is not None:
+                self._write_process.join()
+            logger.info(f"Predictions and manifest saved to {self._output_dir}")
+
+    def _gather_queue_items(self, items: List[QUEUE_ITEM]) -> List[QUEUE_ITEM]:
+        """Gather queue items across distributed ranks, returning only on rank zero."""
+        if not dist_utils.is_distributed():
+            return items
+
+        world_size = dist.get_world_size()
+        object_list: List[List[QUEUE_ITEM]] = [[] for _ in range(world_size)]
+        dist.all_gather_object(object_list, items)
+
+        if self._is_rank_zero:
+            gathered: List[QUEUE_ITEM] = []
+            for rank_items in object_list:
+                gathered.extend(rank_items)
+            return gathered
+
+        return []
 
     def _initialize_write_process(self) -> None:
         self._write_queue = multiprocessing.Queue()

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/core/data/samplers/classification/balanced.py

@@ -1,9 +1,10 @@
 """Random class sampler for data loading."""
 
 from collections import defaultdict
-from typing import Dict, Iterator, List
+from typing import Dict, Iterator, List, Union
 
 import numpy as np
+import torch
 from loguru import logger
 from typing_extensions import override
 
@@ -32,7 +33,7 @@ class BalancedSampler(SamplerWithDataSource[int]):
         """
         self._num_samples = num_samples
         self._replacement = replacement
-        self._class_indices: Dict[int, List[int]] = defaultdict(list)
+        self._class_indices: Dict[Union[int, str], List[int]] = defaultdict(list)
         self._random_generator = np.random.default_rng(seed)
         self._indices: List[int] = []
 
@@ -62,20 +63,31 @@ class BalancedSampler(SamplerWithDataSource[int]):
         super().set_dataset(data_source)
         self._make_indices()
 
+    def _get_class_idx(self, idx):
+        """Load and validate the class index for a given sample index."""
+        if hasattr(self.data_source, "load_target"):
+            target = self.data_source.load_target(idx)  # type: ignore
+        else:
+            _, target, _ = DataSample(*self.data_source[idx])
+
+        if target is None:
+            raise ValueError("The dataset must return non-empty targets.")
+
+        if isinstance(target, str):
+            return target
+
+        if isinstance(target, torch.Tensor):
+            if target.numel() != 1:
+                raise ValueError("The dataset must return a single & scalar target.")
+            return int(target.item())
+
+        raise ValueError("Unsupported target type. Expected str or tensor-like object.")
+
     def _make_indices(self):
         """Samples the indices for each class in the dataset."""
         self._class_indices.clear()
         for idx in tqdm(range(len(self.data_source)), desc="Fetching class indices for sampler"):
-            if hasattr(self.data_source, "load_target"):
-                target = self.data_source.load_target(idx)  # type: ignore
-            else:
-                _, target, _ = DataSample(*self.data_source[idx])
-            if target is None:
-                raise ValueError("The dataset must return non-empty targets.")
-            if target.numel() != 1:
-                raise ValueError("The dataset must return a single & scalar target.")
-
-            class_idx = int(target.item())
+            class_idx = self._get_class_idx(idx)
             self._class_indices[class_idx].append(idx)
 
         if not self._replacement:

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/core/loggers/utils/wandb.py

@@ -5,6 +5,8 @@ from typing import Any, Dict
 
 from loguru import logger
 
+from eva.core.utils import requirements
+
 
 def rename_active_run(name: str) -> None:
     """Renames the current run."""
@@ -12,7 +14,8 @@ def rename_active_run(name: str) -> None:
 
     if wandb.run:
         wandb.run.name = name
-        wandb.run.save()
+        if requirements.below("wandb", "0.21.0"):
+            wandb.run.save()
     else:
         logger.warning("No active wandb run found that could be renamed.")
 

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/core/trainers/trainer.py

@@ -31,6 +31,8 @@ class Trainer(pl_trainer.Trainer):
         default_root_dir: str = "logs",
         n_runs: int = 1,
         checkpoint_type: Literal["best", "last"] = "best",
+        accelerator: str = "auto",
+        devices: int = 1,
         **kwargs: Any,
     ) -> None:
         """Initializes the trainer.
@@ -45,9 +47,17 @@ class Trainer(pl_trainer.Trainer):
             n_runs: The amount of runs (fit and evaluate) to perform in an evaluation session.
             checkpoint_type: Wether to load the "best" or "last" checkpoint saved by the checkpoint
                 callback for evaluations on validation & test sets.
+            accelerator: The accelerator to use for training (e.g. "cpu", "gpu").
+            devices: The number of devices (GPUs) to use for training.
             kwargs: Kew-word arguments of ::class::`lightning.pytorch.Trainer`.
         """
-        super().__init__(*args, default_root_dir=default_root_dir, **kwargs)
+        super().__init__(
+            *args,
+            default_root_dir=default_root_dir,
+            accelerator=accelerator,
+            devices=devices,
+            **kwargs,
+        )
 
         self.checkpoint_type = checkpoint_type
         self.n_runs = n_runs

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/core/utils/__init__.py

@@ -3,5 +3,6 @@
 from eva.core.utils.clone import clone
 from eva.core.utils.memory import to_cpu
 from eva.core.utils.operations import numeric_sort
+from eva.core.utils.paths import home_dir
 
-__all__ = ["clone", "to_cpu", "numeric_sort"]
+__all__ = ["clone", "to_cpu", "numeric_sort", "home_dir"]

kaiko_eva-0.4.1/src/eva/core/utils/distributed.py

@@ -0,0 +1,12 @@
+"""Utility functions for distributed training."""
+
+import torch.distributed as dist
+
+
+def is_distributed() -> bool:
+    """Check if current environment is distributed.
+
+    Returns:
+        bool: True if distributed environment (e.g. multiple gpu processes).
+    """
+    return dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1

kaiko_eva-0.4.1/src/eva/core/utils/paths.py

@@ -0,0 +1,14 @@
+"""Utility functions for handling paths."""
+
+import os
+
+
+def home_dir():
+    """Get eva's home directory for caching."""
+    torch_home = os.path.expanduser(
+        os.getenv(
+            "EVA_HOME",
+            os.path.join("~/.cache", "eva"),
+        )
+    )
+    return torch_home

kaiko_eva-0.4.1/src/eva/core/utils/requirements.py

@@ -0,0 +1,72 @@
+"""Utility functions related to package requirements."""
+
+import importlib
+from typing import Dict
+
+import packaging.version
+
+
+def fetch_version(name: str) -> str | None:
+    """Fetch the installed version of a package.
+
+    Args:
+        name: The name of the package.
+
+    Returns:
+        A string representing the installed version of the package, or None if not found.
+    """
+    try:
+        module = importlib.import_module(name)
+        return getattr(module, "__version__", None)
+    except ImportError:
+        return None
+
+
+def below(name: str, version: str) -> bool:
+    """Check if the installed version of a package is below a certain version.
+
+    Args:
+        name: The name of the package.
+        version: The version to compare against.
+
+    Returns:
+        True if the installed version is below the specified version, False otherwise.
+    """
+    actual = fetch_version(name)
+    if actual:
+        return packaging.version.parse(actual) < packaging.version.parse(version)
+    return False
+
+
+def above_or_equal(name: str, version: str) -> bool:
+    """Check if the installed version of a package is above a certain version.
+
+    Args:
+        name: The name of the package.
+        version: The version to compare against.
+
+    Returns:
+        True if the installed version is above the specified version, False otherwise.
+    """
+    actual = fetch_version(name)
+    if actual:
+        return packaging.version.parse(actual) >= packaging.version.parse(version)
+    return False
+
+
+def check_min_versions(requirements: Dict[str, str]) -> None:
+    """Check installed package versions against requirements dict.
+
+    Args:
+        requirements: A dictionary where keys are package names and
+            values are minimum required versions.
+
+    Raises:
+        ImportError: If any package does not meet the minimum required version.
+    """
+    for package, min_version in requirements.items():
+        if below(package, min_version):
+            raise ImportError(
+                f"Package '{package}' version {fetch_version(package)} does not meet "
+                f"the minimum required version {min_version}."
+            )

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/language/callbacks/writers/prediction.py

@@ -7,11 +7,13 @@ 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
 
@@ -74,10 +76,14 @@ class TextPredictionWriter(callbacks.BasePredictionWriter, abc.ABC):
 
         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._check_if_exists()
+        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()
@@ -105,11 +111,12 @@ class TextPredictionWriter(callbacks.BasePredictionWriter, abc.ABC):
 
         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.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:
@@ -120,26 +127,45 @@ class TextPredictionWriter(callbacks.BasePredictionWriter, abc.ABC):
     @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}")
+        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():
-            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.")
+            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 predictions
+        return output["generated_text"]
 
     def _check_if_exists(self) -> None:
         """Checks if the output directory already exists and if it should be overwritten."""
@@ -150,7 +176,6 @@ class TextPredictionWriter(callbacks.BasePredictionWriter, abc.ABC):
                 "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

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/language/data/datasets/classification/pubmedqa.py

@@ -121,7 +121,7 @@ class PubMedQA(base.TextClassification):
 
     @override
     def validate(self) -> None:
-        if len(self) != self._expected_dataset_lengths[self._split]:
+        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]}, "

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/language/models/modules/language.py

@@ -1,6 +1,6 @@
 """Model module for language models."""
 
-from typing import Any, List
+from typing import Any
 
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from torch import nn
@@ -9,7 +9,7 @@ from typing_extensions import override
 from eva.core.metrics import structs as metrics_lib
 from eva.core.models.modules import module
 from eva.core.models.modules.utils import batch_postprocess
-from eva.language.models.typings import PredictionBatch, TextBatch
+from eva.language.models.typings import ModelOutput, PredictionBatch, TextBatch
 
 
 class LanguageModule(module.ModelModule):
@@ -33,7 +33,7 @@ class LanguageModule(module.ModelModule):
         self.model = model
 
     @override
-    def forward(self, batch: TextBatch, *args: Any, **kwargs: Any) -> List[str]:
+    def forward(self, batch: TextBatch, *args: Any, **kwargs: Any) -> ModelOutput:
         return self.model(batch)
 
     @override
@@ -46,13 +46,14 @@ class LanguageModule(module.ModelModule):
 
     def _batch_step(self, batch: TextBatch) -> STEP_OUTPUT:
         text, targets, metadata = TextBatch(*batch)
-        predictions = self.forward(batch)
+        output = self.forward(batch)
+
         return {
             "inputs": text,
-            "predictions": predictions,
+            "predictions": output.pop("generated_text"),  # type: ignore
             "targets": targets,
             "metadata": metadata,
-        }
+        } | output
 
 
 class OfflineLanguageModule(module.ModelModule):

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/language/models/typings.py

@@ -1,8 +1,9 @@
 """Type definitions for language models."""
 
-from typing import Any, Dict, Generic, List, TypeVar
+from typing import Any, Dict, Generic, List, TypedDict, TypeVar
 
-from typing_extensions import NamedTuple
+import torch
+from typing_extensions import NamedTuple, NotRequired
 
 from eva.language.data.messages import MessageSeries
 
@@ -37,3 +38,19 @@ class PredictionBatch(NamedTuple, Generic[TargetType]):
 
     metadata: Dict[str, Any] | None
     """Additional metadata."""
+
+
+class ModelOutput(TypedDict):
+    """The output batch produced by the model forward pass."""
+
+    generated_text: List[str]
+    """The text generated by the model."""
+
+    input_ids: NotRequired[torch.Tensor | None]
+    """The token ids of the input text."""
+
+    output_ids: NotRequired[torch.Tensor | None]
+    """The token ids of the model output (usually containing both input and prediction)."""
+
+    attention_mask: NotRequired[torch.Tensor | None]
+    """The attention mask for the input tokens."""

{kaiko_eva-0.4.0 → kaiko_eva-0.4.1}/src/eva/language/models/wrappers/base.py

@@ -1,16 +1,16 @@
 """Base class for language model wrappers."""
 
 import abc
-from typing import Any, Callable, List
+from typing import Any, Callable
 
 from typing_extensions import override
 
 from eva.core.models.wrappers import base
 from eva.language.data.messages import ModelSystemMessage
-from eva.language.models.typings import TextBatch
+from eva.language.models.typings import ModelOutput, TextBatch
 
 
-class LanguageModel(base.BaseModel[TextBatch, List[str]]):
+class LanguageModel(base.BaseModel[TextBatch, ModelOutput]):
     """Base class for language models.
 
     Classes that inherit from this should implement the following methods:
@@ -36,7 +36,7 @@ class LanguageModel(base.BaseModel[TextBatch, List[str]]):
         self.system_message = ModelSystemMessage(content=system_prompt) if system_prompt else None
 
     @override
-    def forward(self, batch: TextBatch) -> List[str]:
+    def forward(self, batch: TextBatch) -> ModelOutput:
         """Forward pass of the model."""
         inputs = self.format_inputs(batch)
         return super().forward(inputs)