ClashAI-300k / inference.py
MiguelLabs's picture
Create inference.py
3a10c0f verified
Raw
History Blame Contribute Delete
6.98 kB
"""
run_clashai_300k.py
Script para baixar e rodar o modelo ClashAI-300k (ClashMigLabs) do Hugging Face.
Repositório: https://huggingface.co/ClashMigLabs/ClashAI-300k
Requisitos:
pip install torch huggingface_hub tokenizers
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer
REPO_ID = "ClashMigLabs/ClashAI-300k"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# ------------------------------------------------------------
# 1. BAIXAR ARQUIVOS DO HUGGING FACE
# ------------------------------------------------------------
print("Baixando arquivos do repositório...")
model_path = hf_hub_download(repo_id=REPO_ID, filename="model.pt")
tokenizer_path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer.json")
# ------------------------------------------------------------
# 2. CARREGAR TOKENIZER
# ------------------------------------------------------------
tokenizer = Tokenizer.from_file(tokenizer_path)
PAD_ID = tokenizer.token_to_id("<pad>")
BOS_ID = tokenizer.token_to_id("<bos>")
EOS_ID = tokenizer.token_to_id("<eos>")
SEP_ID = tokenizer.token_to_id("<sep>")
UNK_ID = tokenizer.token_to_id("<unk>")
VOCAB_SIZE = tokenizer.get_vocab_size()
print(f"Vocabulário: {VOCAB_SIZE} tokens")
# ------------------------------------------------------------
# 3. DEFINIÇÃO DA ARQUITETURA (mesma usada no treinamento)
# ------------------------------------------------------------
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=512):
super().__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
self.register_buffer("pe", pe.unsqueeze(0))
def forward(self, x):
return x + self.pe[:, : x.size(1)]
class TinyGPT(nn.Module):
def __init__(self, vocab_size, d_model, n_heads, n_layers, d_ff, max_len, pad_id, dropout=0.1):
super().__init__()
self.d_model = d_model
self.pad_id = pad_id
self.token_emb = nn.Embedding(vocab_size, d_model, padding_idx=pad_id)
self.pos_enc = PositionalEncoding(d_model, max_len)
self.dropout = nn.Dropout(dropout)
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=n_heads,
dim_feedforward=d_ff,
dropout=dropout,
batch_first=True,
activation="gelu",
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
self.ln_f = nn.LayerNorm(d_model)
self.head = nn.Linear(d_model, vocab_size, bias=False)
self.head.weight = self.token_emb.weight
def generate_causal_mask(self, sz, device):
return torch.triu(torch.ones(sz, sz, device=device), diagonal=1).bool()
def forward(self, x):
seq_len = x.size(1)
pad_mask = (x == self.pad_id)
causal_mask = self.generate_causal_mask(seq_len, x.device)
h = self.token_emb(x) * math.sqrt(self.d_model)
h = self.pos_enc(h)
h = self.dropout(h)
h = self.transformer(h, mask=causal_mask, src_key_padding_mask=pad_mask)
h = self.ln_f(h)
return self.head(h)
# ------------------------------------------------------------
# 4. CARREGAR CHECKPOINT
# ------------------------------------------------------------
print("Carregando checkpoint...")
checkpoint = torch.load(model_path, map_location=device, weights_only=False)
# Tenta extrair hiperparâmetros do checkpoint; usa defaults compatíveis com ~300k params
# caso o checkpoint só contenha o state_dict puro.
if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
state_dict = checkpoint["model_state_dict"]
d_model = checkpoint.get("d_model", 64)
n_heads = checkpoint.get("n_heads", 4)
n_layers = checkpoint.get("n_layers", 2)
d_ff = checkpoint.get("d_ff", 128)
max_len = checkpoint.get("max_len", 64)
vocab_size = checkpoint.get("vocab_size", VOCAB_SIZE)
pad_id = checkpoint.get("pad_id", PAD_ID)
else:
# checkpoint é o state_dict cru
state_dict = checkpoint
d_model = 64
n_heads = 4
n_layers = 2
d_ff = 128
max_len = 64
vocab_size = VOCAB_SIZE
pad_id = PAD_ID
model = TinyGPT(
vocab_size=vocab_size,
d_model=d_model,
n_heads=n_heads,
n_layers=n_layers,
d_ff=d_ff,
max_len=max_len,
pad_id=pad_id,
)
model.load_state_dict(state_dict, strict=False)
model.to(device)
model.eval()
n_params = sum(p.numel() for p in model.parameters())
print(f"Modelo carregado. Parâmetros: {n_params:,}")
# ------------------------------------------------------------
# 5. FUNÇÃO DE INFERÊNCIA
# ------------------------------------------------------------
@torch.no_grad()
def responder(pergunta, max_new_tokens=40, temperature=0.7, top_k=10):
q_ids = tokenizer.encode(pergunta.lower()).ids
ids = [BOS_ID] + q_ids + [SEP_ID]
input_ids = torch.tensor([ids], dtype=torch.long, device=device)
for _ in range(max_new_tokens):
logits = model(input_ids)
next_token_logits = logits[0, -1, :] / temperature
top_values, top_indices = torch.topk(next_token_logits, min(top_k, next_token_logits.size(-1)))
probs = F.softmax(top_values, dim=-1)
next_token = top_indices[torch.multinomial(probs, 1)].item()
if next_token == EOS_ID:
break
input_ids = torch.cat([input_ids, torch.tensor([[next_token]], device=device)], dim=1)
if input_ids.size(1) >= max_len:
break
generated_ids = input_ids[0].tolist()
sep_pos = generated_ids.index(SEP_ID) if SEP_ID in generated_ids else 0
answer_ids = generated_ids[sep_pos + 1:]
answer_ids = [i for i in answer_ids if i not in (BOS_ID, EOS_ID, PAD_ID, SEP_ID)]
return tokenizer.decode(answer_ids)
# ------------------------------------------------------------
# 6. LOOP INTERATIVO / TESTE
# ------------------------------------------------------------
if __name__ == "__main__":
perguntas_teste = [
"como conseguir ouro?",
"como derrotar um corredor?",
"o que é elixir duplo?",
]
print("\n=== Teste rápido ===")
for p in perguntas_teste:
r = responder(p)
print(f"Pergunta: {p}")
print(f"Resposta: {r}")
print("-" * 60)
print("\n=== Modo interativo (digite 'sair' para encerrar) ===")
while True:
pergunta = input("\nVocê: ").strip()
if pergunta.lower() in ("sair", "exit", "quit"):
break
if not pergunta:
continue
resposta = responder(pergunta)
print(f"ClashAI: {resposta}")