import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel

model_dir = "./nova1_model"
tokenizer = GPT2Tokenizer.from_pretrained(model_dir)
model = GPT2LMHeadModel.from_pretrained(model_dir)
model.eval()

def generate_text(prompt, max_length=100):
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            do_sample=True,
            top_p=0.95,
            top_k=50,
            temperature=0.8,
            pad_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

if __name__ == "__main__":
    print("Chat NOVA ready. Type your prompt (type 'exit' to quit):")
    while True:
        prompt = input(">> ")
        if prompt.lower() == "exit":
            break
        response = generate_text(prompt)
        print(response)