| """ |
| 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") |
|
|
| |
| |
| |
| 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") |
|
|
| |
| |
| |
| 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") |
|
|
| |
| |
| |
| 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) |
|
|
|
|
| |
| |
| |
| print("Carregando checkpoint...") |
| checkpoint = torch.load(model_path, map_location=device, weights_only=False) |
|
|
| |
| |
| 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: |
| |
| 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:,}") |
|
|
| |
| |
| |
| @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) |
|
|
|
|
| |
| |
| |
| 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}") |