[feat] Add json parser and enable parser config (#19)

* Add json parser and enable parser config

Signed-off-by: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>

Author: arekay-nv

src/inference_endpoint/commands/benchmark.py

@@ -176,7 +176,9 @@ async def run_benchmark_command(args: argparse.Namespace) -> None:
         # ===== YAML MODE - Load from config file =====
         config_path = args.config  # Required by argparse
         try:
-            effective_config = ConfigLoader.load_yaml(Path(config_path))
+            effective_config: BenchmarkConfig = ConfigLoader.load_yaml(
+                Path(config_path)
+            )
 
             # Only auxiliary params allowed (output)
             mode_str = getattr(args, "mode", None)
@@ -203,7 +205,9 @@ async def run_benchmark_command(args: argparse.Namespace) -> None:
     elif benchmark_mode_str in ("offline", "online"):
         # ===== CLI MODE - Build config from CLI params =====
         benchmark_mode = TestType(benchmark_mode_str)  # TestType values are lowercase
-        effective_config = _build_config_from_cli(args, benchmark_mode_str)
+        effective_config: BenchmarkConfig = _build_config_from_cli(
+            args, benchmark_mode_str
+        )
         test_mode = (
             TestMode(args.mode) if getattr(args, "mode", None) else TestMode.PERF
         )
@@ -264,7 +268,7 @@ def _build_config_from_cli(
                 name=args.dataset.stem,
                 type=DatasetType.PERFORMANCE,
                 path=str(args.dataset),
-                format="pkl",  # Will be inferred by DataLoaderFactory
+                format=None,  # Will be inferred by DataLoaderFactory
             )
         ],
         settings=Settings(
@@ -289,6 +293,7 @@ def _build_config_from_cli(
             ),
         ),
         model_params=ModelParams(
+            name=args.model,
             temperature=0.7,
             max_new_tokens=args.max_output_tokens if args.max_output_tokens else 1024,
             osl_distribution=OSLDistribution(
@@ -327,8 +332,7 @@ def _get_dataset_path(args: argparse.Namespace, config: BenchmarkConfig) -> Path
     2. Validate all dataset paths exist
     3. Support dataset interleaving strategies
     """
-    # Priority: CLI args > config
-    if args.dataset:
+    if hasattr(args, "dataset") and args.dataset:
         dataset_path = Path(args.dataset)
     else:
         # TODO: Multi-dataset - currently just picks single dataset
@@ -431,6 +435,8 @@ def _run_benchmark(
     model_name = getattr(args, "model", None)
     if not model_name and config.submission_ref:
         model_name = config.submission_ref.model
+    if not model_name and config.model_params.name:
+        model_name = config.model_params.name
 
     if model_name:
         try:
@@ -476,17 +482,17 @@ def _run_benchmark(
             logger.info("Streaming: disabled (auto, offline mode)")
 
     try:
-        # Create loader using factory
-        def parser(x):
-            return {
-                "prompt": x.text_input,
-                "output": x.ref_output,
-                "model": model_name,
-                "stream": enable_streaming,  # Enable streaming only for online mode
-            }
+        if any(d.parser for d in config.datasets):
+            key_maps = [d.parser for d in config.datasets]
+        else:
+            key_maps = None
+        logger.info(f"Parser key maps: {key_maps}")
 
         dataloader = DataLoaderFactory.create_loader(
-            dataset_path, format=dataset_format, parser=parser
+            dataset_path,
+            format=dataset_format,
+            key_maps=key_maps,
+            metadata={"model": model_name, "stream": enable_streaming},
         )
         dataloader.load()
         logger.info(f"Loaded {dataloader.num_samples()} samples")

src/inference_endpoint/config/schema.py

@@ -135,6 +135,7 @@ class OSLDistribution(BaseModel):
 class ModelParams(BaseModel):
     """Model generation parameters."""
 
+    name: str | None = None
     temperature: float = 0.7
     top_k: int | None = None
     top_p: float | None = None
@@ -179,9 +180,10 @@ class Dataset(BaseModel):
     name: str
     type: DatasetType
     path: str
-    format: str = "pkl"
+    format: str | None = None
     samples: int | None = None
     eval_method: EvalMethod | None = None
+    parser: dict | None = None
 
 
 class RuntimeConfig(BaseModel):

src/inference_endpoint/config/templates/offline_template.yaml

@@ -4,21 +4,24 @@ version: "1.0"
 type: "offline"
 
 model_params:
+  name: "meta-llama/Llama-3.1-8B-Instruct"
   temperature: 0.7
   top_p: 0.9
   max_new_tokens: 1024
 
 datasets:
   - name: "perf-test"
     type: "performance"
-    path: "datasets/openorca.pkl"
+    path: "tests/datasets/dummy_1k.pkl"
     format: "pkl"
     samples: 1000
+    parser:
+      prompt: "text_input"
 
 settings:
   runtime:
-    min_duration_ms: 600000 # 10 minutes
-    max_duration_ms: 1800000 # 30 minutes
+    min_duration_ms: 60000 # 1 minutes
+    max_duration_ms: 180000 # 3 minutes
     scheduler_random_seed: 42 # For Poisson/distribution sampling
     dataloader_random_seed: 42 # For dataset shuffling
 

src/inference_endpoint/config/templates/online_template.yaml

@@ -4,21 +4,23 @@ version: "1.0"
 type: "online"
 
 model_params:
+  name: "meta-llama/Llama-3.1-8B-Instruct"
   temperature: 0.7
   top_p: 0.9
   max_new_tokens: 1024
 
 datasets:
   - name: "latency-test"
     type: "performance"
-    path: "datasets/queries.pkl"
-    format: "pkl"
+    path: "cnn_dailymail_train.json"
     samples: 500
+    parser:
+      prompt: "article"
 
 settings:
   runtime:
-    min_duration_ms: 600000 # 10 minutes
-    max_duration_ms: 1800000 # 30 minutes
+    min_duration_ms: 60000 # 1 minutes
+    max_duration_ms: 180000 # 3 minutes
     scheduler_random_seed: 42 # For Poisson/distribution sampling
     dataloader_random_seed: 42 # For dataset shuffling
 

src/inference_endpoint/dataset_manager/dataloader.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import pickle
 from abc import ABC, abstractmethod
 from collections.abc import Callable
@@ -158,7 +159,7 @@ def __init__(
         self.parser = parser
         self.logger = getLogger(__name__)
         if parser is None:
-
+            # TODO : remove this default implementation
             def extract_text_input(row):
                 return row.text_input
 
@@ -228,7 +229,9 @@ def load_sample(self, index: int) -> Any:
         Loads a sample from the data.
         """
         assert self.loaded, "Data is not loaded. Call load() to load the data."
-        return self.parser(self.data.iloc[index])
+        x = self.parser(self.data.iloc[index])
+        self.logger.debug(f"Loaded sample from pickle file at {index} with keys: {x}")
+        return x
 
     def get_column_names(self):
         return self.data.columns
@@ -289,3 +292,34 @@ def __init__(self, file_path, parser: Callable[[Any], Any] = None):
             parser (Callable[[Any], Any], optional): Callable to parse individual data samples. If not provided, defaults to the parent class's parsing mechanism.
         """
         super().__init__(file_path, parser=parser)
+
+
+class JsonlReader(DataLoader):
+    def __init__(
+        self,
+        file_path,
+        parser: Callable[[Any], Any] = None,
+        metadata: dict | None = None,
+    ):
+        if parser is None:
+            # TODO: Implement a parser interface where yaml files specify the fields to pars
+            def default_parser(x):
+                # Use cnn/daily mail dataset as an example for now.
+                return {"prompt": x["article"]} | metadata
+
+            parser = default_parser
+        super().__init__()
+        self.file_path = file_path
+        self.data = []
+        self.parser = parser
+
+    def load(self):
+        with open(self.file_path) as file:
+            for line in file:
+                self.data.append(self.parser(json.loads(line)))
+
+    def load_sample(self, index: int) -> Any:
+        return self.data[index]
+
+    def num_samples(self):
+        return len(self.data)

src/inference_endpoint/dataset_manager/factory.py

@@ -19,12 +19,12 @@
 """
 
 import logging
-from collections.abc import Callable
 from pathlib import Path
 
 from .dataloader import (
     DataLoader,
     HFDataLoader,
+    JsonlReader,
     PickleReader,
 )
 
@@ -44,15 +44,16 @@ class DataLoaderFactory:
     def create_loader(
         dataset_path: Path | str,
         format: str = "pkl",
-        parser: Callable | None = None,
+        key_maps: list[dict[str, str]] | None = None,
+        metadata: dict | None = None,
         **kwargs,
     ) -> DataLoader:
         """Create appropriate dataset loader based on format.
 
         Args:
             dataset_path: Path to dataset file or directory
             format: Dataset format ("pkl", "jsonl", "hf")
-            parser: Optional parser function for data transformation
+            key_maps: Dictionary of key mappings for the parser
             **kwargs: Additional arguments for specific loaders
 
         Returns:
@@ -61,27 +62,22 @@ def create_loader(
         Raises:
             ValueError: If format is unsupported
         """
-        format = format.lower()
-
-        if format == "pkl" or format == "pickle":
-            # Pickle format - use DeepSeekR1ChatCompletionDataLoader
-            if parser is None:
-                # Default parser for chat completion format
-                def default_parser(x):
-                    return {"prompt": x.text_input, "output": x.ref_output}
+        if key_maps is None:
+            # Assume that the `prompt` key already exists in the dataset
+            key_maps = [{"prompt": "text_input"}]
 
-                parser = default_parser
+        def parser(x):
+            # TODO : handle the entire key_maps list
+            return {k: x[v] for k, v in key_maps[0].items()} | (metadata or {})
 
+        format = format.lower()
+        if format == "pkl" or format == "pickle":
             logger.info(f"Creating pickle dataset loader for {dataset_path}")
             return PickleReader(dataset_path, parser=parser)
 
         elif format == "jsonl" or format == "json":
             # JSON Lines format
-            # TODO: Implement JSONLDataLoader
-            logger.error("JSONL format not yet implemented")
-            raise NotImplementedError(
-                "JSONL dataset format not yet supported. " "Supported formats: pkl, hf"
-            )
+            return JsonlReader(dataset_path, parser=parser, metadata=metadata)
 
         elif format == "hf" or format == "huggingface":
             # HuggingFace dataset

src/inference_endpoint/openai/openai_adapter.py

@@ -29,7 +29,6 @@
     ModelIdsShared,
     Object7,
     ReasoningEffort,
-    Role,
     Role5,
     Role6,
     ServiceTier,
@@ -71,14 +70,14 @@ def to_openai_request(query: Query) -> CreateChatCompletionRequest:
 
         request = CreateChatCompletionRequest(
             model=ModelIdsShared(query.data.get("model", "no-model-name")),
-            # service_tier=ServiceTier.auto,
             reasoning_effort=ReasoningEffort.medium,
             messages=[
-                {
-                    "role": Role.assistant.value,
-                    "content": "You are a helpful assistant.",
-                },
                 {"role": Role5.user.value, "content": query.data["prompt"]},
+                # TODO remove this once we have a way to handle the assistant message
+                # {
+                #     "role": Role.assistant.value,
+                #     "content": "You are a helpful assistant.",
+                # },
             ],
             stream=query.data.get("stream", False),
             max_completion_tokens=query.data.get("max_completion_tokens", 100),

tests/unit/openai/test_openai_types.py

@@ -61,7 +61,8 @@ def test_create_chat_completion_request_from_query(self):
         )
 
         messages = query.model_dump(mode="json")["messages"]
-        assert len(messages) == 2, f"Expected 2 messages, got {len(messages)}"
+        assert len(messages) == 1, f"Expected 1 messages, got {len(messages)}"
+        # TODO : cleanup this once we have a way to handle the assistant message
         for message in messages:
             assert message["role"] in [
                 "assistant",
@@ -73,10 +74,10 @@ def test_create_chat_completion_request_from_query(self):
             assert (
                 message["name"] is None
             ), f"Expected name to be None, got {message['name']}"
-            if message["role"] == "assistant":
-                assert message["content"] == "You are a helpful assistant."
-            else:
+            if message["role"] == "user":
                 assert message["content"] == "Test prompt"
+            # TODO : cleanup this once we have a way to handle the assistant message
+            #     assert message["content"] == "You are a helpful assistant."
 
     def test_create_chat_completion_response_from_query_result(self):
         message_content = "You are a helpful assistant."

tests/unit/test_core_types.py

@@ -41,8 +41,9 @@ def test_query_creation(self) -> None:
         query = OpenAIAdapter.to_openai_request(
             Query(id="test-123", data=payload)
         ).model_dump(mode="json")
-        assert query["messages"][0]["content"] == "You are a helpful assistant."
-        assert query["messages"][1]["content"] == "Test prompt"
+        assert query["messages"][0]["content"] == "Test prompt"
+        # TODO : remove this once we have a way to handle the assistant message
+        # assert query["messages"][1]["content"] == "You are a helpful assistant."
         assert query["model"] == "test-model"
         assert query["max_completion_tokens"] == 100
         assert query["temperature"] == 0.7  # default value
@@ -60,8 +61,9 @@ def test_query_store_load(self) -> None:
         query_loaded = OpenAIAdapter.to_openai_request(
             Query(id="test-123", data=payload)
         )
-        assert query_loaded.messages[0].root.content == "You are a helpful assistant."
-        assert query_loaded.messages[1].root.content == payload["prompt"]
+        assert query_loaded.messages[0].root.content == payload["prompt"]
+        # TODO : remove this once we have a way to handle the assistant message
+        # assert query_loaded.messages[1].root.content == "You are a helpful assistant."
         assert query_loaded.model.root == payload["model"]
         assert query_loaded.max_completion_tokens == payload["max_completion_tokens"]
         assert query_loaded.temperature == payload["temperature"]