import json

import mlx.core as mx
import mlx.nn as nn

from src.tokenizer import BPETokenizer
from src.transformer import MLXGPT as GPT


def build_model(config_path: str = "config.json") -> GPT:
    with open(config_path) as f:
        config = json.load(f)

    return GPT(
        vocab_size=config["vocab_size"],
        embedding_dim=config["embedding_dim"],
        num_layers=config["num_layers"],
        num_heads=config["num_heads"],
        num_query_groups=config["num_query_groups"],
        hidden_dim=config["hidden_dim"],
        max_seq_len=config["max_seq_len"],
        dropout=0.0,
    )


def main():
    device = mx.default_device()
    print(f"Running on {device}")

    model = build_model()
    model.load_weights("model.safetensors")
    model.eval()
    mx.eval(model.parameters())

    tokenizer = BPETokenizer.load("tokenizer")
    prompt = "The future of machine learning"
    prompt_ids = mx.array([tokenizer.encode(prompt)], dtype=mx.int32)

    generated = model.generate(
        prompt_ids,
        max_new_tokens=40,
        temperature=0.8,
        top_k=50,
    )
    mx.eval(generated)
    print(tokenizer.decode(generated[0].tolist()))


if __name__ == "__main__":
    main()