#!/usr/bin/env python3
"""Gera respostas usando exclusivamente o GPT treinado do zero."""

from __future__ import annotations

import argparse
from pathlib import Path

import torch
from tokenizers import Tokenizer

from treinar_gpt_zero import GPTZero, gerar


def carregar(checkpoint_path: Path, device: torch.device):
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    tokenizer_path = Path(checkpoint["tokenizer"])
    if not tokenizer_path.is_absolute():
        tokenizer_path = checkpoint_path.parent / tokenizer_path
    tokenizer = Tokenizer.from_file(str(tokenizer_path))
    modelo = GPTZero(
        checkpoint["vocab_size"],
        checkpoint["block_size"],
        checkpoint["d_model"],
        checkpoint["nhead"],
        checkpoint["num_layers"],
        checkpoint["dropout"],
    ).to(device)
    modelo.load_state_dict(checkpoint["model_state"])
    modelo.eval()
    bos_id = tokenizer.token_to_id("<bos>")
    eos_id = tokenizer.token_to_id("<eos>")
    return modelo, tokenizer, checkpoint, bos_id, eos_id


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path, default=Path("modelo_gpt_zero_stage2.pt"))
    parser.add_argument("--pergunta", default="Olá! Tudo bem? Quem é você?")
    parser.add_argument("--system", default="Você é um assistente útil e claro.")
    parser.add_argument("--tokens", type=int, default=180)
    parser.add_argument("--temperature", type=float, default=0.45)
    parser.add_argument("--top-k", type=int, default=40)
    parser.add_argument("--top-p", type=float, default=0.9)
    parser.add_argument("--repetition-penalty", type=float, default=1.15)
    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
    args = parser.parse_args()
    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)
    modelo, tokenizer, checkpoint, bos_id, eos_id = carregar(args.checkpoint, device)
    prompt = f"<|system|>{args.system}<|end|><|user|>{args.pergunta}<|end|><|assistant|>"
    resposta = gerar(modelo, tokenizer, prompt, device, args.tokens, args.temperature, args.top_k, args.top_p, bos_id, eos_id, args.repetition_penalty)
    print(f"Checkpoint: {args.checkpoint}")
    print(f"Pré-treinado: {checkpoint.get('pretrained')}")
    print(f"Pergunta: {args.pergunta}")
    print(f"Modelo: {resposta}")


if __name__ == "__main__":
    main()