Merge f57a3e2619 into 7050ed204b

run.py

@@ -32,7 +32,6 @@ 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,
@@ -58,7 +57,6 @@ def main():
             model_axis="model",
         ),
     )
-  
     inference_runner = InferenceRunner(
         pad_sizes=(1024,),
         runner=ModelRunner(
@@ -75,14 +73,11 @@ 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