Added cmd args fo top_k and refactored code for readability

Author: ajkv-google

src/maxtext/trainers/post_train/distillation/save_top_k_teacher_logits.py

@@ -6,12 +6,16 @@
 import os
 import pickle
 from typing import Sequence
+import argparse
+import sys
 import tensorflow as tf
 
-from absl import app
 import jax
 import numpy as np
+import functools
+from itertools import islice
 
+from absl import app
 from MaxText import pyconfig
 from maxtext.utils import model_creation_utils
 from maxtext.input_pipeline import input_pipeline_interface
@@ -28,7 +32,7 @@ def get_top_k_logits(logits: jax.Array, k: int):
   return top_k_values, top_k_indices
 
 
-def generate_and_save_data(config):
+def generate_and_save_data(config, k_val):
   """Generates top-k logits from the teacher model and saves them to an ArrayRecord file"""
   devices = jax.devices()
   devices_array = maxtext_utils.create_device_mesh(config, devices)
@@ -45,7 +49,6 @@ def generate_and_save_data(config):
   if config.run_name:
     output_dir = os.path.join(output_dir, config.run_name)
 
-  # create final GCS path and local temp file path
   final_gcs_file = os.path.join(output_dir, "teacher_top_k.array_record")
   local_temp_file = "/tmp/teacher_top_k.array_record"
 
@@ -55,35 +58,26 @@ def generate_and_save_data(config):
     writer = array_record_module.ArrayRecordWriter(local_temp_file, "group_size:1000")
 
   max_logging.log(f"Starting Top-K generation loop for {config.steps} steps...")
-  for step, batch in enumerate(train_iter):
-    if step >= config.steps:
-      break
-
+  for step, batch in enumerate(islice(train_iter, config.steps)):
     tokens = batch["inputs"]
-
     logits = teacher_model(
         decoder_input_tokens=tokens,
         decoder_positions=batch["inputs_position"],
         enable_dropout=False,
     )
 
-    # determine top-k size and extract top-k logits and indices
-    k_val = config.decode_sampling_top_k if config.decode_sampling_top_k > 0 else 128
     top_k_vals, top_k_idx = get_top_k_logits(logits, k=k_val)
 
     gathered_vals = multihost_utils.process_allgather(top_k_vals, tiled=True)
     gathered_idx = multihost_utils.process_allgather(top_k_idx, tiled=True)
     gathered_tokens = multihost_utils.process_allgather(tokens, tiled=True)
 
     optional_keys = ["inputs_position", "inputs_segmentation", "targets_segmentation", "targets"]
-
-    gathered_optionals = {}
-    for key in optional_keys:
-      if key in batch:
-        gathered_optionals[key] = multihost_utils.process_allgather(batch[key], tiled=True)
+    gathered_optionals = {
+        key: multihost_utils.process_allgather(batch[key], tiled=True) for key in optional_keys if key in batch
+    }
 
     if process_index == 0:
-      # Writing the gathered tokens, top-k logits, and top-k indices to the ArrayRecord file
       record_dict = {
           "tokens": np.array(gathered_tokens),
           "top_k_logits": np.array(gathered_vals),
@@ -107,21 +101,31 @@ def generate_and_save_data(config):
   if not tf.io.gfile.exists(output_dir):
     tf.io.gfile.makedirs(output_dir)
 
-  # Upload the local file to GCS and remove the local temp file
   tf.io.gfile.copy(local_temp_file, final_gcs_file, overwrite=True)
   os.remove(local_temp_file)
   max_logging.log("Upload complete")
 
 
-def main(argv: Sequence[str]):
+def main(argv: Sequence[str], local_args):
   # Initialize the global configuration
   global_config = pyconfig.initialize(argv)
   teacher_overrides = global_config.teacher_overrides
   teacher_argv = [argv[0], argv[1]]
   teacher_config = pyconfig.initialize(teacher_argv, **teacher_overrides)
 
-  generate_and_save_data(teacher_config)
+  generate_and_save_data(teacher_config, local_args.top_k)
 
 
 if __name__ == "__main__":
-  app.run(main)
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      "--top_k",
+      type=int,
+      required=False,
+      default=128,
+      help="Top K value for logits.",
+  )
+  local_arg, remaining_args = parser.parse_known_args()
+
+  main_wrapper = functools.partial(main, local_args=local_arg)
+  app.run(main_wrapper, argv=[sys.argv[0]] + remaining_args)