Merge 6ed2d78bea into 7050ed204b

It looks like there were some formatting issues in the code. I've taken the liberty of re-formatting it to be more readable.

run.py

@@ -32,6 +32,7 @@ def validate_checkpoint(path, expected_hash):
 def main():
   # Validate checkpoint integrity
   validate_checkpoint(CKPT_PATH, CKPT_HASH)
+  
   grok_1_model = LanguageModelConfig(
     vocab_size=128 * 1024,
     pad_token=0, 
@@ -57,6 +58,7 @@ def main():
       model_axis="model",
     ),
   )
+  
   inference_runner = InferenceRunner(
     pad_sizes=(1024,),
     runner=ModelRunner(
@@ -73,11 +75,14 @@ def main():
     local_mesh_config=(1, 8),
     between_hosts_config=(1, 1),
   )
+  
   inference_runner.initialize()
+  
   gen = inference_runner.run()
 
   inp = "The answer to life the universe and everything is of course"
   print(f"Output for prompt: {inp}", sample_from_model(gen, inp, max_len=100, temperature=0.01))
+
   # Add authentication
   @app.route("/inference")
   @auth.login_required