import torch
from model import load_qed, Qdevice, encode, sp

model = load_qed("qed-base-300m.pt")

@torch.no_grad()
def generate(model, prompt, max_tokens=100):
    tokens = encode(prompt)
    x = torch.tensor(tokens,dtype=torch.long,device=Qdevice)[None, :]
    for _ in range(max_tokens):
        logits = model(x)
        next_token = torch.argmax(logits[:, -1],dim=-1)
        x = torch.cat([x, next_token[:, None]],dim=1)
    return sp.decode(x[0].tolist())


if __name__ == "__main__":

    while True:
        prompt = input("\nPrompt: ")
        propmt = f"""User:
{prompt}
Assistant:"""
        
        if prompt == "exit":
            break

        print(generate(model, prompt))