savev

Author: Dylan Huang

eval_protocol/data_loader/inline_data_loader.py

@@ -0,0 +1,63 @@
+from collections.abc import Sequence
+
+from eval_protocol.data_loader.models import (
+    DataLoaderContext,
+    DataLoaderResult,
+    DataLoaderVariant,
+    EvaluationDataLoader,
+)
+from eval_protocol.models import EvaluationRow, Message
+from eval_protocol.pytest.types import InputMessagesParam
+
+
+class InlineDataLoader(EvaluationDataLoader):
+    """Data loader for inline ``EvaluationRow`` or message payloads."""
+
+    rows: Sequence[EvaluationRow] | None = None
+    messages: Sequence[InputMessagesParam] | None = None
+    variant_id: str = "inline"
+    description: str | None = None
+
+    def __post_init__(self) -> None:
+        if self.rows is None and self.messages is None:
+            raise ValueError("InlineDataLoader requires rows or messages to be provided")
+
+    def variants(self) -> Sequence[DataLoaderVariant]:
+        def _load(ctx: DataLoaderContext) -> DataLoaderResult:
+            resolved_rows: list[EvaluationRow] = []
+            if self.rows is not None:
+                resolved_rows.extend(row.model_copy(deep=True) for row in self.rows)
+            if self.messages is not None:
+                for dataset_messages in self.messages:
+                    row_messages: list[Message] = []
+                    for msg in dataset_messages:
+                        if isinstance(msg, Message):
+                            row_messages.append(msg.model_copy(deep=True))
+                        else:
+                            row_messages.append(Message.model_validate(msg))
+                    resolved_rows.append(EvaluationRow(messages=row_messages))
+
+            if ctx.max_rows is not None:
+                resolved_rows = resolved_rows[: ctx.max_rows]
+
+            metadata = {
+                "data_loader_variant_id": self.variant_id,
+                "data_loader_type": "inline",
+                "row_count": len(resolved_rows),
+            }
+
+            return DataLoaderResult(
+                rows=resolved_rows,
+                source_id=self.variant_id,
+                source_metadata=metadata,
+            )
+
+        description = self.description or self.variant_id
+        return [
+            DataLoaderVariant(
+                id=self.variant_id,
+                description=description,
+                loader=_load,
+                metadata={"type": "inline"},
+            )
+        ]

eval_protocol/data_loader/models.py

@@ -0,0 +1,71 @@
+"""Data loader abstractions"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Callable
+from typing_extensions import Protocol
+
+from pydantic import BaseModel, Field
+
+from eval_protocol.models import EvaluationRow
+from eval_protocol.pytest.types import EvaluationTestMode
+from eval_protocol.dataset_logger.dataset_logger import DatasetLogger
+
+
+class DataLoaderContext(BaseModel):
+    """Context provided to loader variants when materializing data."""
+
+    max_rows: int | None = Field(default=None, ge=1, description="Maximum number of rows to load")
+    preprocess_fn: Callable[[list[EvaluationRow]], list[EvaluationRow]] | None = Field(
+        default=None, description="Optional preprocessing function for evaluation rows"
+    )
+    logger: DatasetLogger = Field(description="Dataset logger for tracking operations")
+    invocation_id: str = Field(description="Unique identifier for this invocation")
+    experiment_id: str = Field(description="Unique identifier for this experiment")
+    mode: EvaluationTestMode = Field(description="The evaluation test mode")
+
+    class Config:
+        arbitrary_types_allowed = True  # For Callable and DatasetLogger types
+
+
+class DataLoaderResult(BaseModel):
+    """Rows and metadata returned by a loader variant."""
+
+    rows: list[EvaluationRow] = Field(description="List of evaluation rows loaded")
+    source_id: str = Field(description="Unique identifier for the data source")
+    source_metadata: dict[str, Any] = Field(
+        default_factory=dict, description="Additional metadata about the data source"
+    )
+    raw_payload: Any | None = Field(default=None, description="Raw payload data if available")
+    preprocessed: bool = Field(default=False, description="Whether the data has been preprocessed")
+
+    class Config:
+        arbitrary_types_allowed = True  # For Any type in raw_payload
+
+
+class DataLoaderVariant(BaseModel):
+    """Single parameterizable variant from a data loader."""
+
+    id: str = Field(description="Unique identifier for this variant")
+    description: str = Field(description="Human-readable description of this variant")
+    loader: Callable[[DataLoaderContext], DataLoaderResult] = Field(
+        description="Function that loads data for this variant"
+    )
+    metadata: dict[str, Any] = Field(default_factory=dict, description="Additional metadata for this variant")
+
+    class Config:
+        arbitrary_types_allowed = True  # For Callable type
+
+    def load(self, ctx: DataLoaderContext) -> DataLoaderResult:
+        """Load a dataset for this variant using the provided context."""
+
+        return self.loader(ctx)
+
+
+class EvaluationDataLoader(Protocol):
+    """Protocol for data loaders that can be consumed by ``evaluation_test``."""
+
+    def variants(self) -> Sequence[DataLoaderVariant]:
+        """Return parameterizable variants emitted by this loader."""
+        ...

eval_protocol/pytest/evaluation_test.py

@@ -11,6 +11,7 @@
 import pytest
 from tqdm import tqdm
 
+from eval_protocol.data_loader.models import EvaluationDataLoader
 from eval_protocol.dataset_logger import default_logger
 from eval_protocol.dataset_logger.dataset_logger import DatasetLogger
 from eval_protocol.human_id import generate_id, num_combinations
@@ -69,6 +70,7 @@ def evaluation_test(
     input_messages: Sequence[list[InputMessagesParam] | None] | None = None,
     input_dataset: Sequence[DatasetPathParam] | None = None,
     input_rows: Sequence[list[EvaluationRow]] | None = None,
+    input_data_loaders: Sequence[EvaluationDataLoader] | EvaluationDataLoader | None = None,
     dataset_adapter: Callable[[list[dict[str, Any]]], Dataset] = default_dataset_adapter,  # pyright: ignore[reportExplicitAny]
     rollout_processor: RolloutProcessor | None = None,
     evaluation_test_kwargs: Sequence[EvaluationInputParam | None] | None = None,
@@ -131,6 +133,7 @@ def evaluation_test(
         input_rows: Pre-constructed EvaluationRow objects to use directly. This is useful
             when you want to provide EvaluationRow objects with custom metadata, input_messages,
             or other fields already populated. Will be passed as "input_dataset" to the test function.
+        input_loaders: Data loaders to use to load the input dataset.
         dataset_adapter: Function to convert the input dataset to a list of
             EvaluationRows. This is useful if you have a custom dataset format.
         completion_params: Generation parameters for the rollout.
@@ -171,6 +174,11 @@ def evaluation_test(
 
     active_logger: DatasetLogger = logger if logger else default_logger
 
+    if input_data_loaders is not None and (
+        input_dataset is not None or input_messages is not None or input_rows is not None
+    ):
+        raise ValueError("data_loaders cannot be combined with input_dataset, input_messages, or input_rows.")
+
     # Optional global overrides via environment for ad-hoc experimentation
     # EP_INPUT_PARAMS_JSON can contain a JSON object that will be deep-merged
     # into input_params (e.g., '{"temperature":0,"extra_body":{"reasoning":{"effort":"low"}}}').

tests/data_loader/test_inline_data_loader.py

@@ -0,0 +1,18 @@
+from eval_protocol.data_loader.inline_data_loader import InlineDataLoader
+from eval_protocol.models import EvaluationRow, Message
+from eval_protocol.pytest import evaluation_test
+from eval_protocol.pytest.default_no_op_rollout_processor import NoOpRolloutProcessor
+
+
+@evaluation_test(
+    data_loaders=InlineDataLoader(
+        messages=[[Message(role="user", content="What is 2 + 2?")]],
+    ),
+)
+def test_inline_data_loader(row: EvaluationRow) -> EvaluationRow:
+    """Inline data loader should feed pre-constructed message bundles."""
+
+    assert row.messages[0].content == "What is 2 + 2?"
+    assert row.input_metadata.dataset_info is not None
+    assert row.input_metadata.dataset_info.get("data_loader_variant_id") == "inline"
+    return row