"""Talk to Sprout in the terminal.

    python chat.py runs/chat/ckpt_final.pt            # the chat model
    python chat.py runs/base/ckpt_final.pt --base     # the base model just continues your text
"""
import argparse

import torch

from model import GPT, Config
from tokenizer import Tokenizer
from train import get_device


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('ckpt')
    ap.add_argument('--data', default='data')
    ap.add_argument('--base', action='store_true', help='plain continuation, no chat template')
    ap.add_argument('--temperature', type=float, default=0.7)
    ap.add_argument('--top-p', type=float, default=0.9)
    ap.add_argument('--max', type=int, default=200)
    args = ap.parse_args()

    device = get_device()
    tok = Tokenizer.load(f'{args.data}/tokenizer.json')
    ck = torch.load(args.ckpt, map_location='cpu')
    model = GPT(Config(**ck['config']))
    model.load_state_dict(ck['model'])
    model.to(device).eval()
    eot, user, assistant, end = (tok[s] for s in ('<|endoftext|>', '<|user|>', '<|assistant|>', '<|end|>'))

    history = [eot]
    print('Sprout is listening. Empty line to quit.')
    while True:
        text = input('you> ').strip()
        if not text:
            break
        if args.base:
            ids = [eot] + tok.encode(text, allow_special=False)
        else:
            history += [user] + tok.encode(text, allow_special=False) + [end, assistant]
            ids = history[-(model.cfg.context - args.max):]
        out = model.generate(torch.tensor([ids], device=device), args.max, temperature=args.temperature,
                             top_p=args.top_p, stop={end, eot})
        new = out[0, len(ids):].tolist()
        reply = [i for i in new if i not in (end, eot)]
        print('sprout>', (text + ' ' if args.base else '') + tok.decode(reply).strip())
        if not args.base:
            history += reply + [end]


if __name__ == '__main__':
    main()