re-formatting it to be more readable run.py

It looks like there were some formatting issues in the code. I've taken the liberty of re-formatting it to be more readable.
This commit is contained in:
Michael G. Inso 2024-03-26 12:24:32 +03:00 committed by GitHub
parent f57a3e2619
commit 6ed2d78bea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

5
run.py
View File

@ -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