Fixed CR comments

Author: TomerG711

src/guidellm/preprocess/dataset.py

@@ -13,13 +13,7 @@
 
 from guidellm.dataset import load_dataset as guidellm_load_dataset
 from guidellm.utils import IntegerRangeSampler, check_load_processor
-
-SUPPORTED_TYPES = {
-    ".json",
-    ".jsonl",
-    ".csv",
-    ".parquet",
-}
+from guidellm.utils.hf_datasets import SUPPORTED_TYPES, save_dataset_to_file
 
 
 class PromptTooShortError(Exception):
@@ -94,6 +88,7 @@ def handle_pad_strategy(
     min_prompt_tokens: int,
     tokenizer: PreTrainedTokenizerBase,
     pad_char: str,
+    pad_multiplier: int = 2,
     **_kwargs,
 ) -> str:
     """
@@ -103,13 +98,18 @@ def handle_pad_strategy(
     :param min_prompt_tokens: Desired minimum token count.
     :param tokenizer: Tokenizer used to count tokens.
     :param pad_char: Character used for padding.
+    :param pad_multiplier: Multiplier for padding character length.
     :return: Padded prompt string.
     """
 
-    while len(tokenizer.encode(current_prompt)) < min_prompt_tokens:
-        current_prompt += pad_char
-    return current_prompt
-
+    tokens = tokenizer.encode(current_prompt)
+    pad_count = 1
+    prompt = current_prompt
+    while len(tokens) < min_prompt_tokens:
+        prompt += pad_char * pad_count
+        tokens = tokenizer.encode(prompt)
+        pad_count *= pad_multiplier
+    return prompt
 
 def handle_error_strategy(
     current_prompt: str,
@@ -221,31 +221,6 @@ def parse_config_file(data: Union[str, Path]) -> "TokensConfig":
         return TokensConfig(**config_dict)
 
 
-def save_dataset_to_file(dataset: Dataset, output_path: Union[str, Path]) -> None:
-    """
-    Saves a HuggingFace Dataset to file in a supported format.
-
-    :param dataset: Dataset to save.
-    :param output_path: Output file path (.json, .jsonl, .csv, .parquet).
-    :raises ValueError: If the file extension is not supported.
-    """
-    output_path = Path(output_path)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    suffix = output_path.suffix.lower()
-
-    if suffix == ".csv":
-        dataset.to_csv(output_path)
-    elif suffix in {".json", ".jsonl"}:
-        dataset.to_json(output_path)
-    elif suffix == ".parquet":
-        dataset.to_parquet(output_path)
-    else:
-        raise ValueError(
-            f"Unsupported file suffix '{suffix}' in output_path'{output_path}'."
-            f" Only {SUPPORTED_TYPES} are supported."
-        )
-
-
 def _validate_output_suffix(output_path: Union[str, Path]) -> None:
     output_path = Path(output_path)
     suffix = output_path.suffix.lower()
@@ -351,8 +326,8 @@ def process_dataset(
         if prompt_text is None:
             continue
 
-        if len(tokenizer.encode(prompt_text)) > target_prompt_len:
-            tokens = tokenizer.encode(prompt_text)
+        tokens = tokenizer.encode(prompt_text)
+        if len(tokens) > target_prompt_len:
             prompt_text = tokenizer.decode(tokens[:target_prompt_len])
 
         processed_prompt = prompt_row.copy()
@@ -370,7 +345,7 @@ def process_dataset(
 
     processed_dataset = Dataset.from_list(processed_prompts)
     save_dataset_to_file(processed_dataset, output_path)
-    logger.info(f"Conversion complete. Dataset saved to: {output_path}")
+    logger.info(f"Conversion completed. Dataset saved to: {output_path}")
 
     if push_to_hub:
         push_dataset_to_hub(hub_dataset_id, processed_dataset)

src/guidellm/utils/__init__.py

@@ -1,4 +1,8 @@
 from .colors import Colors
+from .hf_datasets import (
+    save_dataset_to_file,
+    SUPPORTED_TYPES,
+)
 from .hf_transformers import (
     check_load_processor,
 )
@@ -22,6 +26,8 @@
     "filter_text",
     "is_puncutation",
     "load_text",
+    "save_dataset_to_file",
     "split_text",
     "split_text_list_by_length",
+    "SUPPORTED_TYPES",
 ]

src/guidellm/utils/hf_datasets.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+from typing import Union
+
+from datasets import Dataset
+
+SUPPORTED_TYPES = {
+    ".json",
+    ".jsonl",
+    ".csv",
+    ".parquet",
+}
+
+
+def save_dataset_to_file(dataset: Dataset, output_path: Union[str, Path]) -> None:
+    """
+    Saves a HuggingFace Dataset to file in a supported format.
+
+    :param dataset: Dataset to save.
+    :param output_path: Output file path (.json, .jsonl, .csv, .parquet).
+    :raises ValueError: If the file extension is not supported.
+    """
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    suffix = output_path.suffix.lower()
+
+    if suffix == ".csv":
+        dataset.to_csv(output_path)
+    elif suffix in {".json", ".jsonl"}:
+        dataset.to_json(output_path)
+    elif suffix == ".parquet":
+        dataset.to_parquet(output_path)
+    else:
+        raise ValueError(
+            f"Unsupported file suffix '{suffix}' in output_path'{output_path}'."
+            f" Only {SUPPORTED_TYPES} are supported."
+        )

tests/unit/preprocess/test_dataset.py

@@ -1,5 +1,4 @@
 import os
-from pathlib import Path
 from typing import TYPE_CHECKING
 from unittest.mock import MagicMock, patch
 
@@ -20,7 +19,6 @@
     handle_pad_strategy,
     process_dataset,
     push_dataset_to_hub,
-    save_dataset_to_file,
 )
 
 
@@ -105,7 +103,7 @@ def test_handle_concatenate_strategy_not_enough_prompts(tokenizer_mock):
 @pytest.mark.sanity
 def test_handle_pad_strategy(tokenizer_mock):
     result = handle_pad_strategy("short", 10, tokenizer_mock, "p")
-    assert result == "shortppppp"
+    assert result.startswith("shortppppp")
 
 
 @pytest.mark.sanity
@@ -122,11 +120,11 @@ def test_handle_error_strategy_too_short_prompt(tokenizer_mock):
 
 
 @pytest.mark.smoke
-@patch("guidellm.preprocess.dataset.save_dataset_to_file")
-@patch("guidellm.preprocess.dataset.Dataset")
-@patch("guidellm.preprocess.dataset.guidellm_load_dataset")
-@patch("guidellm.preprocess.dataset.check_load_processor")
-@patch("guidellm.preprocess.dataset.IntegerRangeSampler")
+@patch(f"{process_dataset.__module__}.save_dataset_to_file")
+@patch(f"{process_dataset.__module__}.Dataset")
+@patch(f"{process_dataset.__module__}.guidellm_load_dataset")
+@patch(f"{process_dataset.__module__}.check_load_processor")
+@patch(f"{process_dataset.__module__}.IntegerRangeSampler")
 def test_process_dataset_non_empty(
     mock_sampler,
     mock_check_processor,
@@ -291,81 +289,3 @@ def test_push_dataset_to_hub_error_no_id():
         push_dataset_to_hub(None, mock_dataset)
 
 
-@pytest.mark.regression
-@patch.object(Path, "mkdir")
-def test_save_dataset_to_file_csv(mock_mkdir):
-    mock_dataset = MagicMock(spec=Dataset)
-    output_path = Path("some/path/output.csv")
-    save_dataset_to_file(mock_dataset, output_path)
-    mock_dataset.to_csv.assert_called_once_with(output_path)
-    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
-
-
-@pytest.mark.regression
-@patch.object(Path, "mkdir")
-def test_save_dataset_to_file_csv_capitalized(mock_mkdir):
-    mock_dataset = MagicMock(spec=Dataset)
-    output_path = Path("some/path/output.CSV")
-    save_dataset_to_file(mock_dataset, output_path)
-    mock_dataset.to_csv.assert_called_once_with(output_path)
-    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
-
-
-@pytest.mark.regression
-@patch.object(Path, "mkdir")
-def test_save_dataset_to_file_json(mock_mkdir):
-    mock_dataset = MagicMock(spec=Dataset)
-    output_path = Path("some/path/output.json")
-    save_dataset_to_file(mock_dataset, output_path)
-    mock_dataset.to_json.assert_called_once_with(output_path)
-    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
-
-
-@pytest.mark.regression
-@patch.object(Path, "mkdir")
-def test_save_dataset_to_file_json_capitalized(mock_mkdir):
-    mock_dataset = MagicMock(spec=Dataset)
-    output_path = Path("some/path/output.JSON")
-    save_dataset_to_file(mock_dataset, output_path)
-    mock_dataset.to_json.assert_called_once_with(output_path)
-    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
-
-
-@pytest.mark.regression
-@patch.object(Path, "mkdir")
-def test_save_dataset_to_file_jsonl(mock_mkdir):
-    mock_dataset = MagicMock(spec=Dataset)
-    output_path = Path("some/path/output.jsonl")
-    save_dataset_to_file(mock_dataset, output_path)
-    mock_dataset.to_json.assert_called_once_with(output_path)
-    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
-
-
-@pytest.mark.regression
-@patch.object(Path, "mkdir")
-def test_save_dataset_to_file_jsonl_capitalized(mock_mkdir):
-    mock_dataset = MagicMock(spec=Dataset)
-    output_path = Path("some/path/output.JSONL")
-    save_dataset_to_file(mock_dataset, output_path)
-    mock_dataset.to_json.assert_called_once_with(output_path)
-    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
-
-
-@pytest.mark.regression
-@patch.object(Path, "mkdir")
-def test_save_dataset_to_file_parquet(mock_mkdir):
-    mock_dataset = MagicMock(spec=Dataset)
-    output_path = Path("some/path/output.parquet")
-    save_dataset_to_file(mock_dataset, output_path)
-    mock_dataset.to_parquet.assert_called_once_with(output_path)
-    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
-
-
-@pytest.mark.regression
-@patch.object(Path, "mkdir")
-def test_save_dataset_to_file_unsupported_type(mock_mkdir):
-    mock_dataset = MagicMock(spec=Dataset)
-    output_path = Path("some/path/output.txt")
-    with pytest.raises(ValueError, match=r"Unsupported file suffix '.txt'.*"):
-        save_dataset_to_file(mock_dataset, output_path)
-    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)

tests/unit/utils/__init__.py

@@ -0,0 +1,87 @@
+from pathlib import Path
+from unittest.mock import patch, MagicMock
+
+import pytest
+from datasets import Dataset
+
+from guidellm.utils import save_dataset_to_file
+
+
+@pytest.mark.regression
+@patch.object(Path, "mkdir")
+def test_save_dataset_to_file_csv(mock_mkdir):
+    mock_dataset = MagicMock(spec=Dataset)
+    output_path = Path("some/path/output.csv")
+    save_dataset_to_file(mock_dataset, output_path)
+    mock_dataset.to_csv.assert_called_once_with(output_path)
+    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
+
+
+@pytest.mark.regression
+@patch.object(Path, "mkdir")
+def test_save_dataset_to_file_csv_capitalized(mock_mkdir):
+    mock_dataset = MagicMock(spec=Dataset)
+    output_path = Path("some/path/output.CSV")
+    save_dataset_to_file(mock_dataset, output_path)
+    mock_dataset.to_csv.assert_called_once_with(output_path)
+    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
+
+
+@pytest.mark.regression
+@patch.object(Path, "mkdir")
+def test_save_dataset_to_file_json(mock_mkdir):
+    mock_dataset = MagicMock(spec=Dataset)
+    output_path = Path("some/path/output.json")
+    save_dataset_to_file(mock_dataset, output_path)
+    mock_dataset.to_json.assert_called_once_with(output_path)
+    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
+
+
+@pytest.mark.regression
+@patch.object(Path, "mkdir")
+def test_save_dataset_to_file_json_capitalized(mock_mkdir):
+    mock_dataset = MagicMock(spec=Dataset)
+    output_path = Path("some/path/output.JSON")
+    save_dataset_to_file(mock_dataset, output_path)
+    mock_dataset.to_json.assert_called_once_with(output_path)
+    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
+
+
+@pytest.mark.regression
+@patch.object(Path, "mkdir")
+def test_save_dataset_to_file_jsonl(mock_mkdir):
+    mock_dataset = MagicMock(spec=Dataset)
+    output_path = Path("some/path/output.jsonl")
+    save_dataset_to_file(mock_dataset, output_path)
+    mock_dataset.to_json.assert_called_once_with(output_path)
+    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
+
+
+@pytest.mark.regression
+@patch.object(Path, "mkdir")
+def test_save_dataset_to_file_jsonl_capitalized(mock_mkdir):
+    mock_dataset = MagicMock(spec=Dataset)
+    output_path = Path("some/path/output.JSONL")
+    save_dataset_to_file(mock_dataset, output_path)
+    mock_dataset.to_json.assert_called_once_with(output_path)
+    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
+
+
+@pytest.mark.regression
+@patch.object(Path, "mkdir")
+def test_save_dataset_to_file_parquet(mock_mkdir):
+    mock_dataset = MagicMock(spec=Dataset)
+    output_path = Path("some/path/output.parquet")
+    save_dataset_to_file(mock_dataset, output_path)
+    mock_dataset.to_parquet.assert_called_once_with(output_path)
+    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
+
+
+@pytest.mark.regression
+@patch.object(Path, "mkdir")
+def test_save_dataset_to_file_unsupported_type(mock_mkdir):
+    mock_dataset = MagicMock(spec=Dataset)
+    output_path = Path("some/path/output.txt")
+    with pytest.raises(ValueError, match=r"Unsupported file suffix '.txt'.*"):
+        save_dataset_to_file(mock_dataset, output_path)
+    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)