"""Load AtoMinima from the Hub and generate text. Run: python example.py"""

import atominima

model = atominima.from_pretrained("arthurblg1802/AtoMinima")
print("parameters:", sum(p.numel() for p in model.parameters()))

for prompt in ["the water cycle begins when", "Question: what is 4 + 5\nAnswer:"]:
    print("\nprompt:", repr(prompt))
    print("output:", repr(model.generate(prompt, max_new_tokens=80)))