| 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)) |
