File size: 2,529 Bytes
9cb4eda | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | #!/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()
|