from m31_runtime import load_stage
from generation_m31 import generate

model, tokenizer, config = load_stage("repair")
prompt = "<bos><user>Write a Python function that returns factorial of n.<assistant>"
print(generate(model, tokenizer, prompt, max_new_tokens=256, temperature=0.7, top_p=0.9, repetition_penalty=1.05, max_context=512))