Repository: AI-Hypercomputer/maxtext
Pull request #1537: Generate dataset for distillation
Author: login: SurbhiJainUSC
Labels: name: pull ready
Created: 2025-04-07T23:48:03Z | Merged: 2025-04-28T20:29:54Z

# Description
This PR is part of adding support for distillation. It creates a script to generate dataset by running inference on a pre-trained teacher model. The generated dataset will then be used to fine-tune a student model.

*Notice 1:* Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

*Notice 2:* For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

# Tests
Command to run MaxEngine Server:
```
python3 -m MaxText.maxengine_server MaxText/configs/base.yml \
tokenizer_path=deepseek-ai/DeepSeek-V2-Lite-chat tokenizer_type=huggingface \
load_parameters_path=gs://sjsurbhi-multipod/deepseek2-16-chat/unscanned/0/items \
model_name=deepseek2-16b per_device_batch_size=10 multi_sampling=True ici_tensor_parallelism=4 \
decode_sampling_strategy=weighted scan_layers=False hf_access_token=$HF_TOKEN
```

 Command to generate distillation dataset:
```
python3 -m MaxText.generate_distillation_data \
--tokenizer-path=deepseek-ai/DeepSeek-V2-Lite-chat \
--dataset-path HuggingFaceH4/ultrachat_200k --data-split train_sft --data-columns messages  \
--hf-repo-id deepseek2-16b-dataset --num-generations 2 \
--hf-access-token=$HF_TOKEN --batch-size 64 --num-batches 2 \
--max-output-length 128 --use-chat-template
```

Generated dataset: https://huggingface.co/datasets/sjsurbhi/deepseek2-16b-dataset

# Checklist

Before submitting this PR, please make sure (put X in square brackets):
- [x] I have performed a self-review of my code.
- [x] I have necessary comments in my code, particularly in hard-to-understand areas.
- [x] I have run end-to-end tests tests and provided workload links above if applicable.
- [x] I have made or will make corresponding changes to the doc if needed.



Diff:
diff --git a/MaxText/generate_distillation_data.py b/MaxText/generate_distillation_data.py
new file mode 100644
index 0000000000..faf5a450b9
--- /dev/null
+++ b/MaxText/generate_distillation_data.py
@@ -0,0 +1,243 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+This script executes the data generation step for Response-based Knowledge Distillation.
+Knowledge Distillation is a compression technique that transfers knowledge
+from a larger (teacher) model to a smaller (student) model.
+The script runs inference on a teacher model to create output samples.
+This generated dataset can be used to fine-tune a student model.
+
+Example command:
+  python3 -m MaxText.generate_distillation_data \
+    --dataset-path HuggingFaceH4/ultrachat_200k --data-split train_sft --data-columns messages \
+    --tokenizer-path deepseek-ai/DeepSeek-V2-Lite-chat \
+    --hf-access-token <access token> --hf-repo-id <hf repository id> \
+    --batch-size 1024 --num-batches 100 \
+    --num-generations 2 \
+    --max-output-length 128 --max-target-length 256 --use-chat-template
+
+Running this command executes 100 processing steps.
+In each step, it generates completions for a batch of 40 prompts.
+This results in inference running on 4000 prompts overall, producing 2 samples per prompt.
+Note:
+Make sure to run maxengine server in a new terminal before executing this command. Example command to run maxengine server:
+  python3 -m MaxText.maxengine_server MaxText/configs/base.yml \
+    model_name=deepseek2-16b tokenizer_path=deepseek-ai/DeepSeek-V2-Lite-chat tokenizer_type=huggingface \
+    load_parameters_path=<unscanned checkpoint path> \
+    per_device_batch_size=10 multi_sampling=True ici_tensor_parallelism=4 \
+    decode_sampling_strategy=weighted scan_layers=False
+"""
+
+import argparse
+import asyncio
+import grpc
+import json
+import transformers
+
+from datasets import Dataset
+from huggingface_hub import create_repo, get_full_repo_name, repo_exists, upload_file
+
+from MaxText import max_logging
+from MaxText.input_pipeline import _distillation_data_processing
+
+from jetstream.core.proto import jetstream_pb2
+from jetstream.core.proto import jetstream_pb2_grpc
+from tqdm.asyncio import tqdm
+
+_GRPC_KEEPALIVE_TIMEOUT_MS = 10000
+_GRPC_MAX_ATTEMPTS = 5
+
+
+async def get_request(input_requests):
+  input_requests = iter(input_requests)
+  for request in input_requests:
+    yield request
+
+
+async def send_request(config, request, stub, tokenizer, progress_bar):  # pylint: disable=redefined-outer-name
+  """Sends the request to JetStream server."""
+  prompt = request.prompt
+  prompt_token_ids = request.prompt_token_ids
+  actual_completion = request.actual_completion
+
+  decode_request = jetstream_pb2.DecodeRequest(
+      token_content=jetstream_pb2.DecodeRequest.TokenContent(token_ids=prompt_token_ids),
+      max_tokens=request.max_output_tokens,
+      num_samples=config.num_generations,  # number of responses to generate for each request
+  )
+
+  response = stub.Decode(decode_request)
+  completion_tokens = [[] for _ in range(config.num_generations)]
+  async for resp in response:
+    for idx, sample in enumerate(resp.stream_content.samples):
+      resp_tokens = sample.token_ids
+      completion_tokens[idx].extend(resp_tokens)
+
+  outputs = []
+  for tokens in completion_tokens:
+    completion = tokenizer.decode(tokens, skip_special_tokens=True)
+    outputs.append(
+        {
+            "prompt": [{"role": "user", "content": prompt}],
+            "completion": [{"role": "assistant", "content": completion}],
+            "actual_completion": [{"role": "assistant", "content": actual_completion}],
+        }
+    )
+  progress_bar.update(1)
+  return outputs
+
+
+async def run_inference(config, requests, tokenizer):  # pylint: disable=redefined-outer-name
+  """Asynchronously runs inference on JetStream server."""
+  progress_bar = tqdm(total=len(requests))
+  progress_bar.set_description(f"Running inference on {len(requests)} prompts")
+
+  server_url = f"localhost:{config.jetstream_server_port}"
+  options = []
+  options.append(("grpc.keepalive_timeout_ms", _GRPC_KEEPALIVE_TIMEOUT_MS))
+  options.append(("grpc.enable_retries", 1))
+  service_config_json = json.dumps(
+      {
+          "methodConfig": [
+              {
+                  "name": [{}],
+                  "retryPolicy": {
+                      "maxAttempts": _GRPC_MAX_ATTEMPTS,
+                      "initialBackoff": "0.2s",
+                      "maxBackoff": "1s",
+                      "backoffMultiplier": 2,
+                      "retryableStatusCodes": ["UNAVAILABLE"],
+                  },
+              }
+          ]
+      }
+  )
+  options.append(("grpc.service_config", service_config_json))
+  tasks = []
+  async with grpc.aio.insecure_channel(server_url, options=options) as channel:
+    stub = jetstream_pb2_grpc.OrchestratorStub(channel)
+    async for request in get_request(requests):
+      tasks.append(
+          asyncio.create_task(
+              send_request(
+                  config=config,
+                  request=request,
+                  stub=stub,
+                  tokenizer=tokenizer,
+                  progress_bar=progress_bar,
+              )
+          )
+      )
+    outputs = await asyncio.gather(*tasks)
+  progress_bar.close()
+  return outputs
+
+
+def generate_completions(config, requests, tokenizer):  # pylint: disable=redefined-outer-name
+  """Generates num_generations of completions for each prompt in request."""
+  outputs = asyncio.run(
+      run_inference(
+          config=config,
+          requests=requests,
+          tokenizer=tokenizer,
+      ),
+  )
+  return [output for output_per_prompt_list in outputs for output in output_per_prompt_list]
+
+
+def upload_data_to_hf(distillation_data, batch_num, hf_repo_id, hf_access_token):
+  """Upload dataset to Hugging Face."""
+  full_repo_name = get_full_repo_name(model_id=hf_repo_id, token=hf_access_token)
+  if not repo_exists(repo_id=full_repo_name, repo_type="dataset", token=hf_access_token):
+    max_logging.log("Repository doesn't exist on Hugging Face, creating a new one.")
+    try:
+      repo_url = create_repo(repo_id=hf_repo_id, repo_type="dataset", private=True, token=hf_access_token)
+      max_logging.log(f"Successfully created repository on Hugging Face: {repo_url}.")
+    except Exception as e:  # pylint: disable=broad-except
+      max_logging.log(f"Error in creating repository on Hugging Face: {e}")
+      raise e
+
+  distillation_dataset = Dataset.from_list(distillation_data)
+  max_logging.log(f"Pushing dataset to Hugging Face: https://huggingface.co/datasets/{full_repo_name}")
+  try:
+    parquet_file_name = f"distillation-data-{batch_num}.parquet"
+    distillation_dataset.to_parquet(parquet_file_name)
+    upload_file(
+        repo_id=full_repo_name,
+        repo_type="dataset",
+        path_or_fileobj=parquet_file_name,
+        path_in_repo=f"data/{parquet_file_name}",
+        commit_message=f"Uploading dataset batch number {batch_num}",
+        token=hf_access_token,
+    )
+    max_logging.log(f"Successfully pushed dataset to Hugging Face: https://huggingface.co/datasets/{full_repo_name}")
+  except Exception as e:  # pylint: disable=broad-except
+    max_logging.log(f"Error in pushing dataset to Hugging Face: {e}")
+    raise e
+
+
+def generate_data(config):  # pylint: disable=redefined-outer-name
+  """Generates data for distillation."""
+  dataset = _distillation_data_processing.load_dataset(config)
+
+  tokenizer = transformers.AutoTokenizer.from_pretrained(
+      config.tokenizer_path,
+      token=config.hf_access_token,
+  )
+
+  start_idx = 0
+  distillation_data = []
+  for batch_num in range(config.num_batches):
+    data = dataset[start_idx : start_idx + config.batch_size]
+    start_idx += config.batch_size
+    sampled_dataset = Dataset.from_dict(data)
+    sampled_dataset = _distillation_data_processing.process_dataset(config, sampled_dataset)
+    requests = _distillation_data_processing.filter_dataset(config, sampled_dataset, tokenizer)
+    distillation_data = generate_completions(config, requests, tokenizer)
+    upload_data_to_hf(distillation_data, batch_num, config.hf_repo_id, config.hf_access_token)
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--jetstream-server-port", type=str, default=9000, help="JetStream server port.")
+  parser.add_argument("--dataset-type", type=str, default="huggingface", help="Type of dataset.")
+  parser.add_argument("--dataset-path", type=str, required=True, help="Path to Hugging Face dataset.")
+  parser.add_argument("--data-split", type=str, required=True, help="Subset of data to load, eg. train or test.")
+  parser.add_argument("--data-columns", nargs="+", required=True, help="Columns names that contain relevant data.")
+  parser.add_argument(
+      "--hf-access-token", type=str, required=True, help="Access token used to load a tokenizer from Hugging Face."
+  )
+  parser.add_argument("--tokenizer-path", type=str, required=True, help="Path to Hugging Face tokenizer.")
+  parser.add_argument("--use-chat-template", action="store_true", help="Enable tokenizer to apply a chat template.")
+  parser.add_argument(
+      "--max-output-length", type=int, required=True, help="The maximum completion tokens to generate for a prompt."
+  )
+  parser.add_argument(
+      "--max-target-length", type=int, default=2048, help="The maximum prompt length plus the output completion length."
+  )
+  parser.add_argument(
+      "--num-generations", type=int, required=False, default=1, help="Number of samples to generate per prompt."
+  )
+  parser.add_argument(
+      "--hf-repo-id", type=str, required=True, help="Name of Hugging Face repository to upload generated dataset."
+  )
+  parser.add_argument("--batch-size", type=int, required=True, help="Number of prompts to process in a batch.")
+  parser.add_argument("--num-batches", type=int, required=True, help="Total number of batches of prompts to process.")
+  config, _ = parser.parse_known_args()
+
+  assert (
+      config.max_output_length < config.max_target_length
+  ), "Maximum output length of completion should be less than maximum target length."
+  generate_data(config)
diff --git a/MaxText/input_pipeline/_distillation_data_processing.py b/MaxText/input_pipeline/_distillation_data_processing.py
new file mode 100644
index 0000000000..45c19eb32e
--- /dev/null
+++ b/MaxText/input_pipeline/_distillation_data_processing.py
@@ -0,0 +1,136 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
+  """
+  example = {
+    "messages": [
+      {"role": "user", "content": "prompt_1"},
+      {"role": "assistant", "content": "completion_1"},
+      {"role": "user", "content": "prompt_2"},
+      {"role": "assistant", "content": "completion_2"}
+    ]
+  }
+  map_to_prompt_completion(example) returns:
+    {
+      "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+      "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+    }
+  """
+  messages = example["messages"]
+  example["prompt"] = [message for message in messages if message["role"] == "user"]
+  example["completion"] = [message for message in messages if message["role"] == "assistant"]
+  return example
+
+
+def extract_content(example, data_column_names):
+  """
+  example = {
+    "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+    "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+  }
+  extract_content(example, ["prompt", "completion"]) returns:
+    {
+      "prompt": ["prompt_1", "prompt_2"],
+      "completion": ["completion_1", "completion_2"]
+    }
+  """
+  for column in data_column_names:
+    example[column] = [data["content"] for data in example[column]]
+  return example
+
+
+def process_dataset(config, dataset):  # pylint: disable=redefined-outer-name
+  """Pipeline for preprocessing dataset."""
+  data_column_names = config.data_columns
+  dataset = dataset.select_columns(data_column_names)
+
+  supported_columns = [["prompt", "completion"], ["messages"]]
+  assert any(
+      set(data_column_names) == set(supported) for supported in supported_columns
+  ), f"Dataset column names mismatch. Expected columns to match one of {supported_columns}, but got {data_column_names}"
+  assert _input_pipeline_utils.is_conversational(
+      dataset.features, data_column_names
+  ), "Dataset is not in conversational format."
+
+  # maps "messages" to "prompt" and "completion" pairs
+  if "messages" in data_column_names:
+    dataset = dataset.map(map_to_prompt_completion, remove_columns=data_column_names)
+
+  data_column_names = ["prompt", "completion"]
+  dataset = dataset.map(extract_content, fn_kwargs={"data_column_names": data_column_names})
+
+  return dataset
+
+
+def load_dataset(config):  # pylint: disable=redefined-outer-name
+  """Loads dataset from Hugging Face."""
+  assert config.dataset_type == "huggingface", "Only dataset from Hugging Face is supported."
+
+  return datasets.load_dataset(
+      config.dataset_path,
+      split=config.data_split,
+      token=config.hf_access_token,
+  )
+
+
+def filter_dataset(config, dataset, tokenizer):
+  "Filter out samples from the dataset."
+  filtered_dataset = []
+  for data in dataset:
+    prompt = data["prompt"][0]
+    actual_completion = data["completion"][0]
+
+    max_output_tokens = min(config.max_output_length, len(tokenizer.encode(actual_completion)))
+    if config.use_chat_template:
+      message = [{"role": "user", "content": prompt}]
+      prompt_token_ids = tokenizer.apply_chat_template(message, add_generation_prompt=True, tokenize=True)
+    else:
+      prompt_token_ids = tokenizer.encode(prompt)
+
+    # Filter out long prompt sequences
+    if len(prompt_token_ids) + max_output_tokens > config.max_target_length:
+      continue
+
+    request = InputRequest(prompt, prompt_token_ids, actual_completion, max_output_tokens)
+    filtered_dataset.append(request)
+  if len(filtered_dataset) < len(dataset):
+    max_logging.log(f"Filtering reduced dataset batch from {len(dataset)} to {len(filtered_dataset)} samples.")
+  return filtered_dataset
diff --git a/MaxText/tests/distillation_data_processing_test.py b/MaxText/tests/distillation_data_processing_test.py
new file mode 100644
index 0000000000..e6e28f4ec5
--- /dev/null
+++ b/MaxText/tests/distillation_data_processing_test.py
@@ -0,0 +1,165 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import argparse
+import os
+import subprocess
+import transformers
+import unittest
+
+from datasets import Dataset
+
+from MaxText.globals import PKG_DIR
+from MaxText.input_pipeline import _distillation_data_processing
+
+PROMPT_DATA = [
+    [
+        {"content": "What color is the sky?", "role": "user"},
+        {"content": "Why is the sky blue?", "role": "user"},
+    ],
+    [
+        {"content": "How many days are in a week?", "role": "user"},
+    ],
+]
+
+COMPLETION_DATA = [
+    [
+        {"content": "The sky is blue.", "role": "assistant"},
+        {"content": "The sky appears blue due a phenomemon called Rayleigh scattering.", "role": "assistant"},
+    ],
+    [
+        {"content": "There are 7 days in a week.", "role": "assistant"},
+    ],
+]
+
+MESSAGES_DATA = [
+    [
+        {"content": "What color is the sky?", "role": "user"},
+        {"content": "The sky is blue.", "role": "assistant"},
+        {"content": "Why is the sky blue?", "role": "user"},
+        {"content": "The sky appears blue due a phenomemon called Rayleigh scattering.", "role": "assistant"},
+    ],
+    [
+        {"content": "How many days are in a week?", "role": "user"},
+        {"content": "There are 7 days in a week.", "role": "assistant"},
+    ],
+]
+
+
+def add_arguments_to_parser(parser):
+  parser.add_argument("--data-columns", nargs="+", required=True, help="Columns names that contain relevant data.")
+  parser.add_argument("--use-chat-template", action="store_true", help="Enable tokenizer to apply a chat template.")
+  parser.add_argument(
+      "--max-output-length", type=int, default=8, help="The maximum completion tokens to generate for a prompt."
+  )
+  parser.add_argument(
+      "--max-target-length", type=int, default=16, help="The maximum prompt length plus the output completion length."
+  )
+  return parser
+
+
+class DistillationDataProcessingTest(unittest.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    exit_code = subprocess.call(
+        [
+            "gsutil",
+            "cp",
+            "-r",
+            "gs://maxtext-dataset/hf/llama2-tokenizer",
+            os.path.join(os.path.dirname(PKG_DIR), "assets", ""),
+        ]
+    )
+    if exit_code != 0:
+      raise ValueError(f"Download tokenizer with gsutil cp failed with exit code: {exit_code}")
+
+  def setUp(self):
+    super().setUp()
+    self.tokenizer = transformers.AutoTokenizer.from_pretrained(
+        os.path.join(os.path.dirname(PKG_DIR), "assets", "llama2-tokenizer"),
+    )
+    self.parser = argparse.ArgumentParser()
+    self.parser = add_arguments_to_parser(self.parser)
+
+  def test_data_processing_with_messages(self):
+    config = self.parser.parse_args(["--data-columns", "messages"])
+    dataset = Dataset.from_dict({"messages": MESSAGES_DATA})
+
+    processed_dataset = _distillation_data_processing.process_dataset(config, dataset)
+
+    expected_prompts = [["What color is the sky?", "Why is the sky blue?"], ["How many days are in a week?"]]
+    expected_completions = [
+        ["The sky is blue.", "The sky appears blue due a phenomemon called Rayleigh scattering."],
+        ["There are 7 days in a week."],
+    ]
+
+    self.assertEqual(len(processed_dataset), len(expected_prompts))
+    for idx, data in enumerate(processed_dataset):
+      self.assertEqual(len(data["prompt"]), len(expected_prompts[idx]))
+      for p_idx, prompt in enumerate(expected_prompts[idx]):
+        self.assertEqual(data["prompt"][p_idx], prompt)
+
+      self.assertEqual(len(data["completion"]), len(expected_completions[idx]))
+      for c_idx, completion in enumerate(expected_completions[idx]):
+        self.assertEqual(data["completion"][c_idx], completion)
+
+  def test_data_filtering_with_messages(self):
+    config = self.parser.parse_args(["--data-columns", "messages"])
+    dataset = Dataset.from_dict({"messages": MESSAGES_DATA})
+
+    processed_dataset = _distillation_data_processing.process_dataset(config, dataset)
+    filtered_dataset = _distillation_data_processing.filter_dataset(config, processed_dataset, self.tokenizer)
+
+    self.assertEqual(len(filtered_dataset), 1)
+    self.assertEqual(filtered_dataset[0].prompt, "What color is the sky?")
+    self.assertEqual(filtered_dataset[0].actual_completion, "The sky is blue.")
+
+  def test_data_processing_with_prompt_completion(self):
+    config = self.parser.parse_args(["--data-columns", "prompt", "completion"])
+    dataset = Dataset.from_dict({"prompt": PROMPT_DATA, "completion": COMPLETION_DATA})
+
+    processed_dataset = _distillation_data_processing.process_dataset(config, dataset)
+
+    expected_prompts = [["What color is the sky?", "Why is the sky blue?"], ["How many days are in a week?"]]
+    expected_completions = [
+        ["The sky is blue.", "The sky appears blue due a phenomemon called Rayleigh scattering."],
+        ["There are 7 days in a week."],
+    ]
+
+    self.assertEqual(len(processed_dataset), len(expected_prompts))
+    for idx, data in enumerate(processed_dataset):
+      self.assertEqual(len(data["prompt"]), len(expected_prompts[idx]))
+      for p_idx, prompt in enumerate(expected_prompts[idx]):
+        self.assertEqual(data["prompt"][p_idx], prompt)
+
+      self.assertEqual(len(data["completion"]), len(expected_completions[idx]))
+      for c_idx, completion in enumerate(expected_completions[idx]):
+        self.assertEqual(data["completion"][c_idx], completion)
+
+  def test_data_filtering_with_prompt_completion(self):
+    config = self.parser.parse_args(["--data-columns", "prompt", "completion"])
+    dataset = Dataset.from_dict({"prompt": PROMPT_DATA, "completion": COMPLETION_DATA})
+
+    processed_dataset = _distillation_data_processing.process_dataset(config, dataset)
+    filtered_dataset = _distillation_data_processing.filter_dataset(config, processed_dataset, self.tokenizer)
+
+    self.assertEqual(len(filtered_dataset), 1)
+    self.assertEqual(filtered_dataset[0].prompt, "What color is the sky?")
+    self.assertEqual(filtered_dataset[0].actual_completion, "The sky is blue.")
+
+
+if __name__ == "__main__":
+  unittest.main()


Reviews:
user: login: gobbleturk
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T22:50:18Z
user: login: gobbleturk
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-09T22:50:59Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T21:01:38Z
user: login: gagika
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T17:47:58Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T18:26:25Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T18:39:19Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T18:39:55Z
user: login: bvandermoon
body: Thanks @SurbhiJainUSC for sharing the dataset output. It's not related to this PR, but do you know why some of the completion data has repeated output?

```
[ { "content": " To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To", "role": "assistant" } ]
```

I remember this was an issue previously. @gagika is this the same issue around decode?
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-25T18:34:23Z
user: login: bvandermoon
body: Generally LGTM pending @gagika's comments also. Just a few smaller comments on my side
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-25T18:46:24Z
user: login: gagika
body: Thanks only a few comments,

also, it would be great to have optoin to save generated Dataset in a gs bucket (this can be a follow up PR done with some extra argument for destination)
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-25T21:37:47Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-25T22:16:45Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-25T22:17:10Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-25T22:17:20Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-25T22:18:10Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-25T22:20:25Z
user: login: gagika
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-26T03:54:21Z

Review comments:
diff_hunk: @@ -0,0 +1,193 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Generates dataset for distillation.
path: MaxText/generate_distillation_data.py
user: login: gobbleturk
body: can you explain a bit more in 1-2 sentences what are the inputs outputs - inputs look like a "teacher" checkpoint and outputs is a huggingface dataset?
created_at: 2025-04-09T22:50:18Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,193 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Generates dataset for distillation.
path: MaxText/generate_distillation_data.py
user: login: SurbhiJainUSC
body: Added more information in the comments.
created_at: 2025-04-11T21:01:38Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,234 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+This script executes the data generation step for Response-based Knowledge Distillation.
+Knowledge Distillation is a compression technique that transfers knowledge
+from a larger (teacher) model to a smaller (student) model.
+The script runs inference on a teacher model to create output samples.
+This generated dataset can be used to fine-tune a student model.
+
+Example command:
+  python3 -m MaxText.generate_distillation_data MaxText/configs/base.yml \
+  run_name=<run_name> base_output_directory=<base_output_directory> \
+  model_name=llama2-7b tokenizer_path=meta-llama/Llama-2-7b-hf tokenizer_type=huggingface \
+  dataset_type=hf hf_path='HuggingFaceH4/ultrachat_200k' train_split=train_sft train_data_columns=[messages] \
+  load_parameters_path=<model checkpoint path> \
+  per_device_batch_size=1 max_target_length=1024 max_prefill_predict_length=512 \
+  steps=100 decode_sampling_strategy=greedy ici_tensor_parallelism=4 \
path: MaxText/generate_distillation_data.py
user: login: gagika
body: we can use `greedy` only when 1 sample is sampled, right?

created_at: 2025-04-16T17:17:33Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,234 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""Input pipeline to generate distillation dataset."""
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: gagika
body: Can you provide more comment what datasets and formats are supported, and how to customize them?
created_at: 2025-04-16T17:43:20Z
author_association: COLLABORATOR
path: MaxText/generate_distillation_data.py
user: login: gagika
body: can we lunch a jetstream server and use its endpoint instead?

@mitalisi @vipannalla 
created_at: 2025-04-16T17:44:42Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,234 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+This script executes the data generation step for Response-based Knowledge Distillation.
+Knowledge Distillation is a compression technique that transfers knowledge
+from a larger (teacher) model to a smaller (student) model.
+The script runs inference on a teacher model to create output samples.
+This generated dataset can be used to fine-tune a student model.
+
+Example command:
+  python3 -m MaxText.generate_distillation_data MaxText/configs/base.yml \
+  run_name=<run_name> base_output_directory=<base_output_directory> \
+  model_name=llama2-7b tokenizer_path=meta-llama/Llama-2-7b-hf tokenizer_type=huggingface \
+  dataset_type=hf hf_path='HuggingFaceH4/ultrachat_200k' train_split=train_sft train_data_columns=[messages] \
+  load_parameters_path=<model checkpoint path> \
+  per_device_batch_size=1 max_target_length=1024 max_prefill_predict_length=512 \
+  steps=100 decode_sampling_strategy=greedy ici_tensor_parallelism=4 \
+  --hf-repo-id="<Hugging Face repository id for dataset>" --num_generations=4
+"""
+
+import argparse
+import asyncio
+import grpc
+import sys
+import transformers
+
+from dataclasses import dataclass
+from datasets import Dataset
+
+from MaxText import max_logging
+from MaxText import max_utils
+from MaxText import maxengine_config
+from MaxText import pyconfig
+from MaxText.input_pipeline import _distillation_data_processing
+from MaxText.train import load_next_batch
+
+from jetstream.core.proto import jetstream_pb2
+from jetstream.core.proto import jetstream_pb2_grpc
+from jetstream.core import server_lib
+from tqdm.asyncio import tqdm
+
+_GRPC_KEEPALIVE_TIMEOUT_MS = 10000
+
+
+@dataclass
+class InputRequest:
+  prompt_tokens: str = ""
+  prompt_true_length: int = 0
+
+
+async def get_request(input_requests):
+  input_requests = iter(input_requests)
+  for request in input_requests:
+    yield request
+
+
+async def send_grpc_request(request, server_url, num_generations):
+  """Sends grpc synchronous request since the current grpc server is sync."""
+  options = [("grpc.keepalive_timeout_ms", _GRPC_KEEPALIVE_TIMEOUT_MS)]
+  async with grpc.aio.insecure_channel(server_url, options=options) as channel:
path: MaxText/generate_distillation_data.py
user: login: gagika
body: is grpc.aio.insecure_channel is created and destroyed for every request within send_grpc_request. ?

I think we also need some retry logic and failure handling.
created_at: 2025-04-16T17:47:51Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,234 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+This script executes the data generation step for Response-based Knowledge Distillation.
+Knowledge Distillation is a compression technique that transfers knowledge
+from a larger (teacher) model to a smaller (student) model.
+The script runs inference on a teacher model to create output samples.
+This generated dataset can be used to fine-tune a student model.
+
+Example command:
+  python3 -m MaxText.generate_distillation_data MaxText/configs/base.yml \
+  run_name=<run_name> base_output_directory=<base_output_directory> \
+  model_name=llama2-7b tokenizer_path=meta-llama/Llama-2-7b-hf tokenizer_type=huggingface \
+  dataset_type=hf hf_path='HuggingFaceH4/ultrachat_200k' train_split=train_sft train_data_columns=[messages] \
+  load_parameters_path=<model checkpoint path> \
+  per_device_batch_size=1 max_target_length=1024 max_prefill_predict_length=512 \
+  steps=100 decode_sampling_strategy=greedy ici_tensor_parallelism=4 \
path: MaxText/generate_distillation_data.py
user: login: SurbhiJainUSC
body: Yes, that's right. I will update this example command and can add a comment as well.
created_at: 2025-04-16T18:26:25Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,234 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+This script executes the data generation step for Response-based Knowledge Distillation.
+Knowledge Distillation is a compression technique that transfers knowledge
+from a larger (teacher) model to a smaller (student) model.
+The script runs inference on a teacher model to create output samples.
+This generated dataset can be used to fine-tune a student model.
+
+Example command:
+  python3 -m MaxText.generate_distillation_data MaxText/configs/base.yml \
+  run_name=<run_name> base_output_directory=<base_output_directory> \
+  model_name=llama2-7b tokenizer_path=meta-llama/Llama-2-7b-hf tokenizer_type=huggingface \
+  dataset_type=hf hf_path='HuggingFaceH4/ultrachat_200k' train_split=train_sft train_data_columns=[messages] \
+  load_parameters_path=<model checkpoint path> \
+  per_device_batch_size=1 max_target_length=1024 max_prefill_predict_length=512 \
+  steps=100 decode_sampling_strategy=greedy ici_tensor_parallelism=4 \
+  --hf-repo-id="<Hugging Face repository id for dataset>" --num_generations=4
+"""
+
+import argparse
+import asyncio
+import grpc
+import sys
+import transformers
+
+from dataclasses import dataclass
+from datasets import Dataset
+
+from MaxText import max_logging
+from MaxText import max_utils
+from MaxText import maxengine_config
+from MaxText import pyconfig
+from MaxText.input_pipeline import _distillation_data_processing
+from MaxText.train import load_next_batch
+
+from jetstream.core.proto import jetstream_pb2
+from jetstream.core.proto import jetstream_pb2_grpc
+from jetstream.core import server_lib
+from tqdm.asyncio import tqdm
+
+_GRPC_KEEPALIVE_TIMEOUT_MS = 10000
+
+
+@dataclass
+class InputRequest:
+  prompt_tokens: str = ""
+  prompt_true_length: int = 0
+
+
+async def get_request(input_requests):
+  input_requests = iter(input_requests)
+  for request in input_requests:
+    yield request
+
+
+async def send_grpc_request(request, server_url, num_generations):
+  """Sends grpc synchronous request since the current grpc server is sync."""
+  options = [("grpc.keepalive_timeout_ms", _GRPC_KEEPALIVE_TIMEOUT_MS)]
+  async with grpc.aio.insecure_channel(server_url, options=options) as channel:
path: MaxText/generate_distillation_data.py
user: login: SurbhiJainUSC
body: Yeah, the channel is destroyed after processing single request. Let me see if we can create the channel once and reuse it for all the requests in a batch.

Also, I will add some retry logic. Thanks for pointing that out.
created_at: 2025-04-16T18:39:19Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,234 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""Input pipeline to generate distillation dataset."""
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: SurbhiJainUSC
body: Sure, will add that.
created_at: 2025-04-16T18:39:55Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: bvandermoon
body: Do we need some unit tests for these?
created_at: 2025-04-25T18:32:16Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
+  """
+  example = {
+    "messages": [
+      {"role": "user", "content": "prompt_1"},
+      {"role": "assistant", "content": "completion_1"},
+      {"role": "user", "content": "prompt_2"},
+      {"role": "assistant", "content": "completion_2"}
+    ]
+  }
+  map_to_prompt_completion(example) returns:
+    {
+      "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+      "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+    }
+  """
+  messages = example["messages"]
+  prompt = []
+  completion = []
+  for message in messages:
+    if message["role"] == "user":
+      prompt.append(message)
+    elif message["role"] == "assistant":
+      completion.append(message)
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: bvandermoon
body: Nit: Could do:

```
prompt = [message for message in messages if message["role"] == "user"]
completion = [message for message in messages if message["role"] == "assistant"]
```
created_at: 2025-04-25T18:43:55Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
+  """
+  example = {
+    "messages": [
+      {"role": "user", "content": "prompt_1"},
+      {"role": "assistant", "content": "completion_1"},
+      {"role": "user", "content": "prompt_2"},
+      {"role": "assistant", "content": "completion_2"}
+    ]
+  }
+  map_to_prompt_completion(example) returns:
+    {
+      "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+      "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+    }
+  """
+  messages = example["messages"]
+  prompt = []
+  completion = []
+  for message in messages:
+    if message["role"] == "user":
+      prompt.append(message)
+    elif message["role"] == "assistant":
+      completion.append(message)
+  example["prompt"] = prompt
+  example["completion"] = completion
+  return example
+
+
+def extract_content(example, data_column_names):
+  """
+  example = {
+    "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+    "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+  }
+  extract_content(example, ["prompt", "completion"]) returns:
+    {
+      "prompt": ["prompt_1", "prompt_2"],
+      "completion": ["completion_1", "completion_2"]
+    }
+  """
+  for column in data_column_names:
+    data_list = example[column]
+    content_list = []
+    for data in data_list:
+      content_list.append(data["content"])
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: bvandermoon
body: Nit: Could do:

```
content_list = [data["content"] for data in data_list]
```
created_at: 2025-04-25T18:45:25Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
+  """
+  example = {
+    "messages": [
+      {"role": "user", "content": "prompt_1"},
+      {"role": "assistant", "content": "completion_1"},
+      {"role": "user", "content": "prompt_2"},
+      {"role": "assistant", "content": "completion_2"}
+    ]
+  }
+  map_to_prompt_completion(example) returns:
+    {
+      "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+      "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+    }
+  """
+  messages = example["messages"]
+  prompt = []
+  completion = []
+  for message in messages:
+    if message["role"] == "user":
+      prompt.append(message)
+    elif message["role"] == "assistant":
+      completion.append(message)
+  example["prompt"] = prompt
+  example["completion"] = completion
+  return example
+
+
+def extract_content(example, data_column_names):
+  """
+  example = {
+    "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+    "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+  }
+  extract_content(example, ["prompt", "completion"]) returns:
+    {
+      "prompt": ["prompt_1", "prompt_2"],
+      "completion": ["completion_1", "completion_2"]
+    }
+  """
+  for column in data_column_names:
+    data_list = example[column]
+    content_list = []
+    for data in data_list:
+      content_list.append(data["content"])
+    example[column] = content_list
+  return example
+
+
+def process_dataset(config, dataset):  # pylint: disable=redefined-outer-name
+  """Pipeline for preprocessing dataset."""
+  assert config.tokenizer_type == "huggingface", "Only tokenizer from Hugging Face is supported."
+
+  data_column_names = config.data_columns
+  dataset = dataset.select_columns(data_column_names)
+
+  supported_columns = [["prompt", "completion"], ["messages"]]
+  assert any(
+      set(data_column_names) == set(supported) for supported in supported_columns
+  ), f"Dataset column names mismatch. Expected columns to match one of {supported_columns}, but got {data_column_names}"
+  assert _input_pipeline_utils.is_conversational(
+      dataset.features, data_column_names
+  ), "Dataset is not in conversational format."
+
+  # maps "messages" to "prompt" and "completion" pairs
+  if "messages" in data_column_names:
+    dataset = dataset.map(map_to_prompt_completion, remove_columns=data_column_names)
+
+  data_column_names = ["prompt", "completion"]
+  dataset = dataset.map(extract_content, fn_kwargs={"data_column_names": data_column_names})
+
+  return dataset
+
+
+def load_dataset(config):  # pylint: disable=redefined-outer-name
+  """Loads dataset from Hugging Face."""
+  assert config.dataset_type == "huggingface", "Only dataset from Hugging Face is supported."
+
+  return datasets.load_dataset(
+      config.dataset_path,
+      split=config.data_split,
+      token=config.hf_access_token,
+  )
+
+
+def filter_dataset(config, dataset):
+  "Filter out samples from the dataset."
+  tokenizer = transformers.AutoTokenizer.from_pretrained(
+      config.tokenizer_path,
+      token=config.hf_access_token,
+  )
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: gagika
body: Can this be moved out and passed as an input? so that we don't load it for every batch.
created_at: 2025-04-25T21:30:44Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
+  """
+  example = {
+    "messages": [
+      {"role": "user", "content": "prompt_1"},
+      {"role": "assistant", "content": "completion_1"},
+      {"role": "user", "content": "prompt_2"},
+      {"role": "assistant", "content": "completion_2"}
+    ]
+  }
+  map_to_prompt_completion(example) returns:
+    {
+      "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+      "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+    }
+  """
+  messages = example["messages"]
+  prompt = []
+  completion = []
+  for message in messages:
+    if message["role"] == "user":
+      prompt.append(message)
+    elif message["role"] == "assistant":
+      completion.append(message)
+  example["prompt"] = prompt
+  example["completion"] = completion
+  return example
+
+
+def extract_content(example, data_column_names):
+  """
+  example = {
+    "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+    "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+  }
+  extract_content(example, ["prompt", "completion"]) returns:
+    {
+      "prompt": ["prompt_1", "prompt_2"],
+      "completion": ["completion_1", "completion_2"]
+    }
+  """
+  for column in data_column_names:
+    data_list = example[column]
+    content_list = []
+    for data in data_list:
+      content_list.append(data["content"])
+    example[column] = content_list
+  return example
+
+
+def process_dataset(config, dataset):  # pylint: disable=redefined-outer-name
+  """Pipeline for preprocessing dataset."""
+  assert config.tokenizer_type == "huggingface", "Only tokenizer from Hugging Face is supported."
+
+  data_column_names = config.data_columns
+  dataset = dataset.select_columns(data_column_names)
+
+  supported_columns = [["prompt", "completion"], ["messages"]]
+  assert any(
+      set(data_column_names) == set(supported) for supported in supported_columns
+  ), f"Dataset column names mismatch. Expected columns to match one of {supported_columns}, but got {data_column_names}"
+  assert _input_pipeline_utils.is_conversational(
+      dataset.features, data_column_names
+  ), "Dataset is not in conversational format."
+
+  # maps "messages" to "prompt" and "completion" pairs
+  if "messages" in data_column_names:
+    dataset = dataset.map(map_to_prompt_completion, remove_columns=data_column_names)
+
+  data_column_names = ["prompt", "completion"]
+  dataset = dataset.map(extract_content, fn_kwargs={"data_column_names": data_column_names})
+
+  return dataset
+
+
+def load_dataset(config):  # pylint: disable=redefined-outer-name
+  """Loads dataset from Hugging Face."""
+  assert config.dataset_type == "huggingface", "Only dataset from Hugging Face is supported."
+
+  return datasets.load_dataset(
+      config.dataset_path,
+      split=config.data_split,
+      token=config.hf_access_token,
+  )
+
+
+def filter_dataset(config, dataset):
+  "Filter out samples from the dataset."
+  tokenizer = transformers.AutoTokenizer.from_pretrained(
+      config.tokenizer_path,
+      token=config.hf_access_token,
+  )
+
+  filtered_dataset = []
+  for data in dataset:
+    prompt = data["prompt"][0]
+    actual_completion = data["completion"][0]
+
+    max_output_tokens = min(config.max_output_length, len(tokenizer.encode(actual_completion)))
+    if config.use_chat_template:
+      message = [{"role": "user", "content": prompt}]
+      prompt_token_ids = tokenizer.apply_chat_template(message, add_generation_prompt=True, tokenize=True)
+    else:
+      prompt_token_ids = tokenizer.encode(prompt)
+
+    # Filter out long prompt sequences
+    if len(prompt_token_ids) + max_output_tokens > config.max_target_length:
+      continue
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: gagika
body: could you print a warning here?, so that user knows some prompts are filtered out.
created_at: 2025-04-25T21:32:59Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: SurbhiJainUSC
body: Thanks for pointing that out. I have added unit tests.
created_at: 2025-04-25T22:16:44Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
+  """
+  example = {
+    "messages": [
+      {"role": "user", "content": "prompt_1"},
+      {"role": "assistant", "content": "completion_1"},
+      {"role": "user", "content": "prompt_2"},
+      {"role": "assistant", "content": "completion_2"}
+    ]
+  }
+  map_to_prompt_completion(example) returns:
+    {
+      "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+      "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+    }
+  """
+  messages = example["messages"]
+  prompt = []
+  completion = []
+  for message in messages:
+    if message["role"] == "user":
+      prompt.append(message)
+    elif message["role"] == "assistant":
+      completion.append(message)
+  example["prompt"] = prompt
+  example["completion"] = completion
+  return example
+
+
+def extract_content(example, data_column_names):
+  """
+  example = {
+    "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+    "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+  }
+  extract_content(example, ["prompt", "completion"]) returns:
+    {
+      "prompt": ["prompt_1", "prompt_2"],
+      "completion": ["completion_1", "completion_2"]
+    }
+  """
+  for column in data_column_names:
+    data_list = example[column]
+    content_list = []
+    for data in data_list:
+      content_list.append(data["content"])
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: SurbhiJainUSC
body: Done, thanks!
created_at: 2025-04-25T22:17:10Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
+  """
+  example = {
+    "messages": [
+      {"role": "user", "content": "prompt_1"},
+      {"role": "assistant", "content": "completion_1"},
+      {"role": "user", "content": "prompt_2"},
+      {"role": "assistant", "content": "completion_2"}
+    ]
+  }
+  map_to_prompt_completion(example) returns:
+    {
+      "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+      "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+    }
+  """
+  messages = example["messages"]
+  prompt = []
+  completion = []
+  for message in messages:
+    if message["role"] == "user":
+      prompt.append(message)
+    elif message["role"] == "assistant":
+      completion.append(message)
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: SurbhiJainUSC
body: Done, thanks!
created_at: 2025-04-25T22:17:20Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
+  """
+  example = {
+    "messages": [
+      {"role": "user", "content": "prompt_1"},
+      {"role": "assistant", "content": "completion_1"},
+      {"role": "user", "content": "prompt_2"},
+      {"role": "assistant", "content": "completion_2"}
+    ]
+  }
+  map_to_prompt_completion(example) returns:
+    {
+      "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+      "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+    }
+  """
+  messages = example["messages"]
+  prompt = []
+  completion = []
+  for message in messages:
+    if message["role"] == "user":
+      prompt.append(message)
+    elif message["role"] == "assistant":
+      completion.append(message)
+  example["prompt"] = prompt
+  example["completion"] = completion
+  return example
+
+
+def extract_content(example, data_column_names):
+  """
+  example = {
+    "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+    "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+  }
+  extract_content(example, ["prompt", "completion"]) returns:
+    {
+      "prompt": ["prompt_1", "prompt_2"],
+      "completion": ["completion_1", "completion_2"]
+    }
+  """
+  for column in data_column_names:
+    data_list = example[column]
+    content_list = []
+    for data in data_list:
+      content_list.append(data["content"])
+    example[column] = content_list
+  return example
+
+
+def process_dataset(config, dataset):  # pylint: disable=redefined-outer-name
+  """Pipeline for preprocessing dataset."""
+  assert config.tokenizer_type == "huggingface", "Only tokenizer from Hugging Face is supported."
+
+  data_column_names = config.data_columns
+  dataset = dataset.select_columns(data_column_names)
+
+  supported_columns = [["prompt", "completion"], ["messages"]]
+  assert any(
+      set(data_column_names) == set(supported) for supported in supported_columns
+  ), f"Dataset column names mismatch. Expected columns to match one of {supported_columns}, but got {data_column_names}"
+  assert _input_pipeline_utils.is_conversational(
+      dataset.features, data_column_names
+  ), "Dataset is not in conversational format."
+
+  # maps "messages" to "prompt" and "completion" pairs
+  if "messages" in data_column_names:
+    dataset = dataset.map(map_to_prompt_completion, remove_columns=data_column_names)
+
+  data_column_names = ["prompt", "completion"]
+  dataset = dataset.map(extract_content, fn_kwargs={"data_column_names": data_column_names})
+
+  return dataset
+
+
+def load_dataset(config):  # pylint: disable=redefined-outer-name
+  """Loads dataset from Hugging Face."""
+  assert config.dataset_type == "huggingface", "Only dataset from Hugging Face is supported."
+
+  return datasets.load_dataset(
+      config.dataset_path,
+      split=config.data_split,
+      token=config.hf_access_token,
+  )
+
+
+def filter_dataset(config, dataset):
+  "Filter out samples from the dataset."
+  tokenizer = transformers.AutoTokenizer.from_pretrained(
+      config.tokenizer_path,
+      token=config.hf_access_token,
+  )
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: SurbhiJainUSC
body: Updated, thank you!
created_at: 2025-04-25T22:18:10Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
+  """
+  example = {
+    "messages": [
+      {"role": "user", "content": "prompt_1"},
+      {"role": "assistant", "content": "completion_1"},
+      {"role": "user", "content": "prompt_2"},
+      {"role": "assistant", "content": "completion_2"}
+    ]
+  }
+  map_to_prompt_completion(example) returns:
+    {
+      "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+      "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+    }
+  """
+  messages = example["messages"]
+  prompt = []
+  completion = []
+  for message in messages:
+    if message["role"] == "user":
+      prompt.append(message)
+    elif message["role"] == "assistant":
+      completion.append(message)
+  example["prompt"] = prompt
+  example["completion"] = completion
+  return example
+
+
+def extract_content(example, data_column_names):
+  """
+  example = {
+    "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+    "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+  }
+  extract_content(example, ["prompt", "completion"]) returns:
+    {
+      "prompt": ["prompt_1", "prompt_2"],
+      "completion": ["completion_1", "completion_2"]
+    }
+  """
+  for column in data_column_names:
+    data_list = example[column]
+    content_list = []
+    for data in data_list:
+      content_list.append(data["content"])
+    example[column] = content_list
+  return example
+
+
+def process_dataset(config, dataset):  # pylint: disable=redefined-outer-name
+  """Pipeline for preprocessing dataset."""
+  assert config.tokenizer_type == "huggingface", "Only tokenizer from Hugging Face is supported."
+
+  data_column_names = config.data_columns
+  dataset = dataset.select_columns(data_column_names)
+
+  supported_columns = [["prompt", "completion"], ["messages"]]
+  assert any(
+      set(data_column_names) == set(supported) for supported in supported_columns
+  ), f"Dataset column names mismatch. Expected columns to match one of {supported_columns}, but got {data_column_names}"
+  assert _input_pipeline_utils.is_conversational(
+      dataset.features, data_column_names
+  ), "Dataset is not in conversational format."
+
+  # maps "messages" to "prompt" and "completion" pairs
+  if "messages" in data_column_names:
+    dataset = dataset.map(map_to_prompt_completion, remove_columns=data_column_names)
+
+  data_column_names = ["prompt", "completion"]
+  dataset = dataset.map(extract_content, fn_kwargs={"data_column_names": data_column_names})
+
+  return dataset
+
+
+def load_dataset(config):  # pylint: disable=redefined-outer-name
+  """Loads dataset from Hugging Face."""
+  assert config.dataset_type == "huggingface", "Only dataset from Hugging Face is supported."
+
+  return datasets.load_dataset(
+      config.dataset_path,
+      split=config.data_split,
+      token=config.hf_access_token,
+  )
+
+
+def filter_dataset(config, dataset):
+  "Filter out samples from the dataset."
+  tokenizer = transformers.AutoTokenizer.from_pretrained(
+      config.tokenizer_path,
+      token=config.hf_access_token,
+  )
+
+  filtered_dataset = []
+  for data in dataset:
+    prompt = data["prompt"][0]
+    actual_completion = data["completion"][0]
+
+    max_output_tokens = min(config.max_output_length, len(tokenizer.encode(actual_completion)))
+    if config.use_chat_template:
+      message = [{"role": "user", "content": prompt}]
+      prompt_token_ids = tokenizer.apply_chat_template(message, add_generation_prompt=True, tokenize=True)
+    else:
+      prompt_token_ids = tokenizer.encode(prompt)
+
+    # Filter out long prompt sequences
+    if len(prompt_token_ids) + max_output_tokens > config.max_target_length:
+      continue
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: SurbhiJainUSC
body: I am wondering if that would flood the logs with a lot of warnings. To avoid that, I have added a final log ([here](https://github.com/AI-Hypercomputer/maxtext/pull/1537/files#diff-604e00b4eac91b196918ad7ffcffd394e2dbd4fd8b03bd5c823f53af604ddab5R135)) to convey that some prompts will be filtered out from the batch.
created_at: 2025-04-25T22:20:24Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,153 @@
+#  Copyright 2025 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Input pipeline to generate knowledge distillation dataset from conversational dataset.
+The conversational dataset should conform to one of the two schemas:
+1. Contains a `messages` column: Typically holding a list of message
+   (e.g., [{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}, ...]).
+2. Contains `prompt` and `completion` columns: Separating the input
+    query [{'role': 'user', 'content': '...'}] from the target output [{'role': 'assistant', 'content': '...'}].
+"""
+
+import datasets
+import transformers
+
+from dataclasses import dataclass, field
+from typing import List
+from MaxText import max_logging
+from MaxText.input_pipeline import _input_pipeline_utils
+
+
+@dataclass
+class InputRequest:
+  prompt: str = ""
+  prompt_token_ids: List[int] = field(default_factory=list)
+  actual_completion: str = ""
+  max_output_tokens: int = 0
+
+
+def map_to_prompt_completion(example):
+  """
+  example = {
+    "messages": [
+      {"role": "user", "content": "prompt_1"},
+      {"role": "assistant", "content": "completion_1"},
+      {"role": "user", "content": "prompt_2"},
+      {"role": "assistant", "content": "completion_2"}
+    ]
+  }
+  map_to_prompt_completion(example) returns:
+    {
+      "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+      "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+    }
+  """
+  messages = example["messages"]
+  prompt = []
+  completion = []
+  for message in messages:
+    if message["role"] == "user":
+      prompt.append(message)
+    elif message["role"] == "assistant":
+      completion.append(message)
+  example["prompt"] = prompt
+  example["completion"] = completion
+  return example
+
+
+def extract_content(example, data_column_names):
+  """
+  example = {
+    "prompt": [{"role": "user", "content": "prompt_1"}, {"role": "user", "content": "prompt_2"}],
+    "completion": [{"role": "assistant", "content": "completion_1"}, {"role": "assistant", "content": "completion_2"}]
+  }
+  extract_content(example, ["prompt", "completion"]) returns:
+    {
+      "prompt": ["prompt_1", "prompt_2"],
+      "completion": ["completion_1", "completion_2"]
+    }
+  """
+  for column in data_column_names:
+    data_list = example[column]
+    content_list = []
+    for data in data_list:
+      content_list.append(data["content"])
+    example[column] = content_list
+  return example
+
+
+def process_dataset(config, dataset):  # pylint: disable=redefined-outer-name
+  """Pipeline for preprocessing dataset."""
+  assert config.tokenizer_type == "huggingface", "Only tokenizer from Hugging Face is supported."
+
+  data_column_names = config.data_columns
+  dataset = dataset.select_columns(data_column_names)
+
+  supported_columns = [["prompt", "completion"], ["messages"]]
+  assert any(
+      set(data_column_names) == set(supported) for supported in supported_columns
+  ), f"Dataset column names mismatch. Expected columns to match one of {supported_columns}, but got {data_column_names}"
+  assert _input_pipeline_utils.is_conversational(
+      dataset.features, data_column_names
+  ), "Dataset is not in conversational format."
+
+  # maps "messages" to "prompt" and "completion" pairs
+  if "messages" in data_column_names:
+    dataset = dataset.map(map_to_prompt_completion, remove_columns=data_column_names)
+
+  data_column_names = ["prompt", "completion"]
+  dataset = dataset.map(extract_content, fn_kwargs={"data_column_names": data_column_names})
+
+  return dataset
+
+
+def load_dataset(config):  # pylint: disable=redefined-outer-name
+  """Loads dataset from Hugging Face."""
+  assert config.dataset_type == "huggingface", "Only dataset from Hugging Face is supported."
+
+  return datasets.load_dataset(
+      config.dataset_path,
+      split=config.data_split,
+      token=config.hf_access_token,
+  )
+
+
+def filter_dataset(config, dataset):
+  "Filter out samples from the dataset."
+  tokenizer = transformers.AutoTokenizer.from_pretrained(
+      config.tokenizer_path,
+      token=config.hf_access_token,
+  )
+
+  filtered_dataset = []
+  for data in dataset:
+    prompt = data["prompt"][0]
+    actual_completion = data["completion"][0]
+
+    max_output_tokens = min(config.max_output_length, len(tokenizer.encode(actual_completion)))
+    if config.use_chat_template:
+      message = [{"role": "user", "content": prompt}]
+      prompt_token_ids = tokenizer.apply_chat_template(message, add_generation_prompt=True, tokenize=True)
+    else:
+      prompt_token_ids = tokenizer.encode(prompt)
+
+    # Filter out long prompt sequences
+    if len(prompt_token_ids) + max_output_tokens > config.max_target_length:
+      continue
path: MaxText/input_pipeline/_distillation_data_processing.py
user: login: gagika
body: I think it's okey, if a lot of prompts are going to be missed because of this setting, it would be good to get notified (so the user can update the params), but the hope is that those are set in such a way that not many are skipped.


created_at: 2025-04-26T03:54:21Z
author_association: COLLABORATOR

PR comments:
user: login: SurbhiJainUSC
created_at: 2025-04-25T22:17:44Z
body: > Thanks only a few comments,
> 
> also, it would be great to have optoin to save generated Dataset in a gs bucket (this can be a follow up PR done with some extra argument for destination)

Sure, I can do that in a follow up PR.
author_association: COLLABORATOR
user: login: SurbhiJainUSC
created_at: 2025-04-25T23:31:05Z
body: > Thanks @SurbhiJainUSC for sharing the dataset output. It's not related to this PR, but do you know why some of the completion data has repeated output?
> 
> ```
> [ { "content": " To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To To", "role": "assistant" } ]
> ```
> 
> I remember this was an issue previously. @gagika is this the same issue around decode?

Yeah, I am aware of this issue. I don't think it is related to this PR, this is what JetStream returns. I will debug this further and try to get better generated results.
author_association: COLLABORATOR