Merge pull request #70 from bespokelabsai/ryanm/batch_size_arg

Add Prompter arg for batch size

examples/distill.py

@@ -3,7 +3,7 @@
 import logging
 
 dataset = load_dataset("allenai/WildChat", split="train")
-dataset = dataset.select(range(3_000))
+dataset = dataset.select(range(300))
 
 # To see more detail about how batches are being processed
 logger = logging.getLogger("bespokelabs.curator")
@@ -20,7 +20,11 @@ def parse_func(row, response):
 
 
 distill_prompter = curator.Prompter(
-    prompt_func=prompt_func, parse_func=parse_func, model_name="gpt-4o-mini", batch=True
+    prompt_func=prompt_func,
+    parse_func=parse_func,
+    model_name="gpt-4o-mini",
+    batch=True,
+    batch_size=100,
 )
 
 distilled_dataset = distill_prompter(dataset)

src/bespokelabs/curator/prompter/prompter.py

@@ -8,6 +8,7 @@
 from datasets import Dataset
 from pydantic import BaseModel
 from xxhash import xxh64
+import logging
 
 from bespokelabs.curator.db import MetadataDB
 from bespokelabs.curator.prompter.prompt_formatter import PromptFormatter
@@ -24,6 +25,8 @@
 
 T = TypeVar("T")
 
+logger = logging.getLogger(__name__)
+
 
 class Prompter:
     """Interface for prompting LLMs."""
@@ -39,6 +42,7 @@ def __init__(
         ] = None,
         response_format: Optional[Type[BaseModel]] = None,
         batch: bool = False,
+        batch_size: Optional[int] = None,
     ):
         """Initialize a Prompter.
 
@@ -50,6 +54,8 @@ def __init__(
                 response object and returns the parsed output
             response_format (Optional[Type[BaseModel]]): A Pydantic model specifying the
                 response format from the LLM.
+            batch (bool): Whether to use batch processing
+            batch_size (Optional[int]): The size of the batch to use, only used if batch is True
         """
         prompt_sig = inspect.signature(prompt_func)
         if len(prompt_sig.parameters) > 1:
@@ -69,8 +75,14 @@ def __init__(
         )
 
         if batch:
-            self._request_processor = OpenAIBatchRequestProcessor(model=model_name)
+            self._request_processor = OpenAIBatchRequestProcessor(
+                model=model_name, batch_size=batch_size
+            )
         else:
+            if batch_size is not None:
+                logger.warning(
+                    f"Prompter argument `batch_size` {batch_size} is ignored because `batch` is False"
+                )
             self._request_processor = OpenAIOnlineRequestProcessor(model=model_name)
 
     def __call__(self, dataset: Optional[Iterable] = None) -> Dataset:

src/bespokelabs/curator/request_processor/openai_online_request_processor.py

@@ -28,12 +28,11 @@
 class OpenAIOnlineRequestProcessor(BaseRequestProcessor):
     def __init__(
         self,
-        batch_size: Optional[int] = None,
         model: str = "gpt-4o-mini",
         api_key: str = os.getenv("OPENAI_API_KEY"),
         url: str = "https://api.openai.com/v1/chat/completions",
     ):
-        super().__init__(batch_size)
+        super().__init__(batch_size=None)
         self.model: str = model
         self.url: str = url
         self.api_key: str = api_key