Fix test

Signed-off-by: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com>

Author: arekay-nv

src/inference_endpoint/dataset_manager/dataloader.py

@@ -163,7 +163,7 @@ def __init__(
         if parser is None:
             # TODO : remove this default implementation
             def extract_text_input(row):
-                return row.text_input
+                return {"prompt": row["text_input"]} | (self.metadata or {})
 
             self.parser = extract_text_input
 
@@ -195,7 +195,7 @@ def load(self, force: bool = False):
                 # note that while we are appending, which is slower, the data may not be sequential and may
                 # have gaps
                 self.text_inputs.append(self.parser(data.to_dict()))
-                
+
         self.loaded = True
 
     def num_samples(self):

tests/conftest.py

@@ -193,10 +193,7 @@ def ds_pickle_reader(ds_pickle_dataset_path):
     """
 
     def parser(row):
-        ret = {}
-        for column in row.index.to_list():
-            ret[column] = row[column]
-        return ret
+        return row
 
     return PickleReader(ds_pickle_dataset_path, parser=parser)
 
@@ -358,7 +355,7 @@ def parser(x):
             Returns:
                 dict: A dictionary with 'prompt' and 'output' keys derived from the input sample.
             """
-            return {"prompt": x.text_input, "output": x.ref_output}
+            return {"prompt": x["text_input"], "output": x["ref_output"]}
 
         self.parser = parser
         data_loader = DeepSeekR1ChatCompletionDataLoader(

tests/integration/test_end_to_end_oracle.py

@@ -121,7 +121,11 @@ async def _run_load_generator_full_run_url(
     url, dataset_path, tmp_path, clean_sample_event_hooks, hf_model_name
 ):
     def parser(x):
-        res = {"prompt": x.text_input, "output": x.ref_output, "model": hf_model_name}
+        res = {
+            "prompt": x["text_input"],
+            "output": x["ref_output"],
+            "model": hf_model_name,
+        }
         return res
 
     dummy_dataloader = DeepSeekR1ChatCompletionDataLoader(dataset_path, parser)
@@ -181,7 +185,11 @@ async def test_load_generator_full_run_mock_http_oracle_server(
     hf_model_name,
 ):
     def parser(x):
-        res = {"prompt": x.text_input, "output": x.ref_output, "model": hf_model_name}
+        res = {
+            "prompt": x["text_input"],
+            "output": x["ref_output"],
+            "model": hf_model_name,
+        }
         return res
 
     dummy_dataloader = DeepSeekR1ChatCompletionDataLoader(

tests/integration/test_server_roundtrip.py

@@ -40,7 +40,7 @@ async def test_ds_chat_completion_data_loader_with_oracle_server(
     """
 
     def parser(x):
-        return {"prompt": x.text_input, "output": x.ref_output}
+        return {"prompt": x["text_input"], "output": x["ref_output"]}
 
     ds_chat_completion_data_loader = DeepSeekR1ChatCompletionDataLoader(
         ds_pickle_dataset_path, parser=parser