mirror of
https://github.com/xai-org/grok-1.git
synced 2024-08-25 16:23:01 +03:00
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:
parent
f57a3e2619
commit
6ed2d78bea
5
run.py
5
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
|
||||
|
Loading…
Reference in New Issue
Block a user