from nanosota_moe import load_nanosota
model, tokenizer = load_nanosota("oscarz511/NanoSOTA-v2-General")

prompt = "If it takes 3 hours to dry 3 shirts, how long for 30 shirts?"
inputs = tokenizer.apply_chat_template([
    {"role": "system", "content": "You are NanoSOTA. Think step-by-step."}, 
    {"role": "user", "content": prompt}
], return_tensors="pt", add_generation_prompt=True).to("cuda")

out = model.generate(**inputs, max_new_tokens=256)
print(tokenizer.decode(out[0]))