"""
Standalone inference for Sid-GPT.

    pip install torch tokenizers
    python generate.py --prompt "Responsible for the mobile app roadmap."
"""
import argparse
import os

import torch
from tokenizers import ByteLevelBPETokenizer

from model import GPT

HERE = os.path.dirname(os.path.abspath(__file__))
TEMPLATE = "Turn this into a strong, quantified PM resume bullet:\n{}"


def load(checkpoint=None, tokenizer_dir=None, device="cpu"):
    tok = ByteLevelBPETokenizer(os.path.join(tokenizer_dir or os.path.join(HERE, "tokenizer"), "vocab.json"),
                                os.path.join(tokenizer_dir or os.path.join(HERE, "tokenizer"), "merges.txt"))
    ckpt = torch.load(checkpoint or os.path.join(HERE, "sft.pt"), map_location=device, weights_only=False)
    model = GPT(ckpt["config"]).to(device)
    model.load_state_dict(ckpt["model"])
    model.eval()
    return model, tok


def rewrite(model, tok, bullet, temperature=0.5, top_k=20, max_new_tokens=70, device="cpu"):
    uid, aid, eot = (tok.token_to_id(t) for t in ("<|user|>", "<|assistant|>", "<|endoftext|>"))
    ids = [uid] + tok.encode(TEMPLATE.format(bullet.strip())).ids + [aid]
    with torch.no_grad():
        out = model.generate(torch.tensor([ids], device=device), max_new_tokens=max_new_tokens,
                             temperature=temperature, top_k=top_k)
    new = out[0, len(ids):].tolist()
    if eot in new:
        new = new[:new.index(eot)]
    return tok.decode(new).strip()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt", required=True)
    ap.add_argument("--temperature", type=float, default=0.5)
    ap.add_argument("--top_k", type=int, default=20)
    ap.add_argument("--max_new_tokens", type=int, default=70)
    args = ap.parse_args()
    model, tok = load()
    print(rewrite(model, tok, args.prompt, args.temperature, args.top_k, args.max_new_tokens))


if __name__ == "__main__":
    main()