removed the need for hardcoding arrayrecord file and read directly from config

Author: ajkv-google

src/maxtext/configs/types.py

@@ -1081,15 +1081,13 @@ class Distillation(BaseModel):
       default_factory=dict,
       description="Overrides specific to the Teacher model (e.g., {'num_query_heads': 64}).",
   )
-  
+
   # --- Offline Distillation Fields ---
   offline_distillation: bool = Field(
-      False, 
-      description="If True, enables offline distillation using pre-generated teacher data."
+      False, description="If True, enables offline distillation using pre-generated teacher data."
   )
   offline_data_dir: Optional[str] = Field(
-      None, 
-      description="GCS or local path to the pre-generated ArrayRecord teacher data."
+      None, description="GCS or local path to the pre-generated ArrayRecord teacher data."
   )
 
   # --- Loss Params ---

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

@@ -18,7 +18,6 @@
 model structures with Tunix's training interfaces.
 """
 
-import os
 import pickle
 import tensorflow as tf
 from array_record.python import array_record_module
@@ -82,11 +81,7 @@ class OfflineArrayRecordIterator:
   """Reads the pre-generated global top-k logits file."""
 
   def __init__(self, data_dir: str, epochs: int = 100):
-    # Check if the user passed a directory or a direct file path
-    if tf.io.gfile.isdir(data_dir):
-      self.filepath = os.path.join(data_dir, "teacher_top_k_global.array_record")
-    else:
-      self.filepath = data_dir
+    self.filepath = data_dir
 
     if not tf.io.gfile.exists(self.filepath):
       raise FileNotFoundError(f"Offline distillation file not found: {self.filepath}")

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

@@ -32,10 +32,6 @@
 3. **Tunix Integration**: We wrap the MaxText models in `TunixMaxTextAdapter` to expose
    a standard interface (call signature) that the Tunix `DistillationTrainer` expects.
 """
-import argparse
-import functools
-import sys
-
 from typing import Sequence, Callable
 from absl import app
 from flax import nnx
@@ -654,4 +650,4 @@ def main(argv: Sequence[str]) -> None:
 
 
 if __name__ == "__main__":
-  app.run(main)
+  app.run(main)