"""Load the exported rurtech.ai MoE and generate text.

  python generate.py --prompt "The rurtech.ai runtime"
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
from safetensors.torch import load_model

from modeling_ecoaco import EcoacoConfig, EcoacoForCausalLM
from tokenizer import ByteBPETokenizer

HERE = Path(__file__).resolve().parent


def load(artifacts: Path):
    cfg_d = json.loads((artifacts / "config.json").read_text())
    for _k in ("model_type","architectures","name"): cfg_d.pop(_k, None)
    cfg = EcoacoConfig(**cfg_d)
    model = EcoacoForCausalLM(cfg)
    load_model(model, str(artifacts / "model.safetensors"))
    model.eval()
    tok = ByteBPETokenizer.load(artifacts / "tokenizer.json")
    return model, tok


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt", default="The rurtech.ai runtime")
    ap.add_argument("--max-new-tokens", type=int, default=60)
    ap.add_argument("--temperature", type=float, default=0.7)
    ap.add_argument("--artifacts", default=str(HERE / "artifacts"))
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, tok = load(Path(args.artifacts))
    model = model.to(device)
    ctx = torch.tensor([tok.encode(args.prompt, add_bos=True)], device=device)
    out = model.generate(ctx, max_new_tokens=args.max_new_tokens,
                         temperature=args.temperature, eos_id=tok.eos_id)
    print(tok.decode(out[0].tolist()))


if __name__ == "__main__":
    main()