""" 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("") BOS_ID = tokenizer.token_to_id("") EOS_ID = tokenizer.token_to_id("") SEP_ID = tokenizer.token_to_id("") UNK_ID = tokenizer.token_to_id("") 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}")