Implement item type

Author: sjmonson

src/guidellm/preprocess/item.py

@@ -0,0 +1,60 @@
+from collections.abc import Sequence
+from typing import Generic, Optional, TypeVar, Union
+
+from pydantic import Field
+
+from guidellm.objects.pydantic import StandardBaseModel
+
+PromptT = TypeVar("PromptT")
+
+
+class Item(StandardBaseModel, Generic[PromptT]):
+    """
+    Represents a single item in a dataset, containing a prompt and its associated metadata.
+    """
+
+    value: PromptT = Field(
+        description="The prompt text or data for the item.",
+        examples=[
+            "What is the capital of France?",
+            "Explain quantum computing in simple terms.",
+        ],
+    )
+    prompt_tokens: Optional[int] = Field(
+        default=None, gt=0, description="Number of tokens in the prompt"
+    )
+    output_tokens: Optional[int] = Field(
+        default=None, gt=0, description="Number of tokens in the output"
+    )
+
+
+class ItemList(Sequence[Item[PromptT]]):
+    """
+    Represents a list of items, each containing a prompt and its metadata.
+    """
+
+    def __init__(self, *items: Item[PromptT], shared_prefix: Optional[PromptT] = None):
+        self.shared_prefix: Optional[PromptT] = shared_prefix
+        self._items: list[Item[PromptT]] = list(items)
+
+    def __getitem__(self, key) -> Union[Item[PromptT], Sequence[Item[PromptT]]]:
+        return self._items[key]
+
+    def __len__(self) -> int:
+        return len(self._items)
+
+    @classmethod
+    def from_lists(
+        cls,
+        prompts: list[PromptT],
+        prompts_tokens: list[Optional[int]],
+        outputs_tokens: list[Optional[int]],
+    ) -> "ItemList":
+        return cls(
+            *[
+                Item(value=prompt, output_tokens=in_t, prompt_tokens=out_t)
+                for prompt, in_t, out_t in zip(
+                    prompts, prompts_tokens, outputs_tokens, strict=True
+                )
+            ]
+        )

src/guidellm/request/loader.py

@@ -11,10 +11,9 @@
 from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
 from transformers import PreTrainedTokenizerBase  # type: ignore[import]
 
-from guidellm.config import settings
 from guidellm.dataset import ColumnInputTypes, load_dataset
 from guidellm.objects import StandardBaseModel
-from guidellm.request.request import GenerationRequest
+from guidellm.preprocess.item import ItemList
 from guidellm.request.session import GenerativeRequestSession
 
 __all__ = [
@@ -107,20 +106,13 @@ def __init__(
         self._preserved_iter = None
 
     def __iter__(self) -> Iterator[GenerativeRequestSession]:
-        turns = 1
-
-        data_iter = self._create_requests()
-        while requests := [i for i, _ in zip(data_iter, range(turns))]:
-            yield GenerativeRequestSession(requests)
-
-    def _create_requests(self) -> Iterator[GenerationRequest]:
         scope_create_count = 0
 
         while (dataset_iter := self._get_dataset_iter(scope_create_count)) is not None:
             scope_create_count += 1
 
             for item in dataset_iter:
-                yield self._create_request(item)
+                yield GenerativeRequestSession(self._create_items(item))
 
             self._preserved_iter = None
 
@@ -268,25 +260,17 @@ def _get_dataset_iter(
 
         return dataset_iter
 
-    def _create_request(self, item: dict[str, Any]) -> GenerationRequest:
-        prompt_tokens = (
-            item[self.column_mappings["prompt_tokens_count_column"]]
+    def _create_items(self, item: dict[str, Any]) -> ItemList:
+        prompts = list(item[self.column_mappings["prompt_column"]])
+        prompt_tokens: list[Optional[int]] = (
+            list(item[self.column_mappings["prompt_tokens_count_column"]])
             if "prompt_tokens_count_column" in self.column_mappings
-            else None
+            else [None]
         )
-        output_tokens = (
-            item[self.column_mappings["output_tokens_count_column"]]
+        output_tokens: list[Optional[int]] = (
+            list(item[self.column_mappings["output_tokens_count_column"]])
             if "output_tokens_count_column" in self.column_mappings
-            else None
+            else [None]
         )
 
-        return GenerationRequest(
-            request_type=settings.preferred_route,
-            content=item[self.column_mappings["prompt_column"]],
-            stats=(
-                {"prompt_tokens": prompt_tokens} if prompt_tokens is not None else {}
-            ),
-            constraints=(
-                {"output_tokens": output_tokens} if output_tokens is not None else {}
-            ),
-        )
+        return ItemList.from_lists(prompts, prompt_tokens, output_tokens)

src/guidellm/request/session.py

@@ -1,16 +1,16 @@
 import itertools
 from abc import ABC, abstractmethod
-from typing import Generic, TypeVar
+from collections.abc import Sequence
+from typing import Generic
 
 from guidellm.backend.response import ResponseSummary
+from guidellm.config import settings
+from guidellm.preprocess.item import Item, ItemList
 from guidellm.request.request import GenerationRequest
+from guidellm.request.types import RequestT, ResponseT
 
 __all__ = ["GenerativeRequestSession", "RequestSession"]
 
-# TODO: Replace with specific types that implement needed features
-RequestT = TypeVar("RequestT")
-ResponseT = TypeVar("ResponseT")
-
 
 class RequestSession(ABC, Generic[RequestT, ResponseT]):
     @abstractmethod
@@ -30,46 +30,61 @@ def push_response(self, response: ResponseT) -> None: ...
     def complete(self) -> bool: ...
 
 
-# FIXME: Bad implementation. Can only handle string requests
 class GenerativeRequestSession(RequestSession[GenerationRequest, ResponseSummary]):
-    def __init__(self, prompts: list[GenerationRequest]) -> None:
-        if not prompts:
+    def __init__(self, items: ItemList) -> None:
+        if len(items) < 1:
             raise ValueError("Prompts cannot be empty")
 
-        self.prompts = prompts
-        self.responses: list[str] = []
+        self.prompts: Sequence[Item] = items
+        self.responses: list[Item] = []
 
     def __len__(self) -> int:
         return len(self.prompts)
 
     def get_next_request(self) -> GenerationRequest:
         completed_responses = len(self.responses)
-        base_request = self.prompts[completed_responses].model_copy(deep=True)
-        base_request.content = "".join(
+
+        # FIXME: Can only handle string requests
+        content = "".join(
             itertools.chain.from_iterable(
-                zip((x.content for x in self.prompts), self.responses + [""])
+                (x.value, y.value)
+                for x, y in zip(self.prompts, self.responses + [Item(value="")])
             )
         )
-        base_request.stats["prompt_tokens"] = sum(
-            x.stats["prompt_tokens"] for x in self.prompts[: completed_responses + 1]
+
+        prev_prompt_tokens = sum(
+            (x.prompt_tokens or 0) + (x.output_tokens or 0) for x in self.responses
         )
-        base_request.constraints["output_tokens"] = sum(
-            x.constraints["output_tokens"]
-            for x in self.prompts[: completed_responses + 1]
+        prompt_tokens = (
+            self.prompts[completed_responses].prompt_tokens or 0
+        ) + prev_prompt_tokens
+
+        output_tokens = self.prompts[completed_responses].output_tokens
+
+        return GenerationRequest(
+            request_type=settings.preferred_route,
+            content=content,
+            stats=(
+                {"prompt_tokens": prompt_tokens} if prompt_tokens is not None else {}
+            ),
+            constraints=(
+                {"output_tokens": output_tokens} if output_tokens is not None else {}
+            ),
         )
 
-        return base_request
-
     def get_next_delay(self) -> float:
         return 0.0
 
     def push_response(self, response: ResponseSummary) -> None:
         if len(self.responses) < len(self.prompts):
-            if response.response_output_tokens is not None:
-                self.prompts[len(self.responses)].constraints["output_tokens"] = (
-                    response.response_output_tokens
-                )
-            self.responses.append(response.value)
+            resp = Item(
+                value=response.value,
+                prompt_tokens=response.response_prompt_tokens
+                or response.request_prompt_tokens,
+                output_tokens=response.response_output_tokens
+                or response.request_output_tokens,
+            )
+            self.responses.append(resp)
         else:
             raise ValueError("Response list full")