from aggpt14 import AgGPT14
from training_data import corpus

if __name__ == "__main__":
    model = AgGPT14(corpus, order=3, seed=None)

    prompt = "What is your favorite color?"
    print(f"User: {prompt}")
    response = model.ask(prompt)
    print(f"AI:   {response}")


    prompt = "hello, how are you?"
    print(f"\nPrompt: {prompt}")
    response = model.ask(prompt, max_tokens=20, temperature=0.5, top_k=5, text_world_model=True)
    print(f"Response: {response}")