Spaces:
Paused
Paused
| import os, json, base64, re, time | |
| import numpy as np | |
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| CONFIG = { | |
| "corpus_file": "Recognizement_Corpus.txt", | |
| "context_length": 12, | |
| "embed_dim": 16, | |
| "num_heads": 2, | |
| "num_layers": 1, | |
| "ffn_hidden": 48, | |
| "batch_size": 128, | |
| "max_vocab": 700, | |
| "learning_rate": 0.015, | |
| "focal_gamma": 0.5, | |
| "loss_norm": True, | |
| "max_epochs": 600, | |
| "target_loss": 0.10, | |
| } | |
| PAD, UNK, EOS = "<pad>", "<unk>", "<eos>" | |
| PAD_ID, UNK_ID, EOS_ID = 0, 1, 2 | |
| def tokenize(text): | |
| text = text.replace("<pad>", " __PAD__ ").replace("<unk>", " __UNK__ ").replace("<eos>", " __EOS__ ") | |
| text = text.lower() | |
| text = re.sub(r"([.,!?;:])", r" \1 ", text) | |
| return text.replace("__PAD__", "<pad>").replace("__UNK__", "<unk>").replace("__EOS__", "<eos>").split() | |
| def build_vocab(tokens, max_vocab=None): | |
| from collections import Counter | |
| freq = Counter(tokens) | |
| vocab = {PAD: 0, UNK: 1, EOS: 2} | |
| for tok, _ in freq.most_common(): | |
| if tok in vocab: continue | |
| if max_vocab and len(vocab) >= max_vocab: break | |
| vocab[tok] = len(vocab) | |
| return vocab | |
| def build_dataset(lines, vocab, ctx): | |
| X_all, Y_all, M_all = [], [], [] | |
| for line in lines: | |
| toks = tokenize(line) + [EOS] | |
| ids = [vocab.get(t, UNK_ID) for t in toks][:ctx+1] | |
| real_len = len(ids) | |
| padded = ids + [PAD_ID] * (ctx + 1 - len(ids)) | |
| X_all.append(padded[:-1]); Y_all.append(padded[1:]) | |
| M_all.append([1.0 if i < real_len - 1 else 0.0 for i in range(ctx)]) | |
| return (torch.tensor(X_all, dtype=torch.long), | |
| torch.tensor(Y_all, dtype=torch.long), | |
| torch.tensor(M_all, dtype=torch.float32)) | |
| class MHSA(nn.Module): | |
| def __init__(self, D, H, ctx): | |
| super().__init__() | |
| self.D, self.H, self.hd = D, H, D // H | |
| self.Wq = nn.Linear(D, D); self.Wk = nn.Linear(D, D) | |
| self.Wv = nn.Linear(D, D); self.Wo = nn.Linear(D, D) | |
| self.register_buffer("mask", torch.triu(torch.full((ctx, ctx), -1e9), diagonal=1)) | |
| def forward(self, x): | |
| B, T, D = x.shape | |
| Q = self.Wq(x).view(B, T, self.H, self.hd).transpose(1, 2) | |
| K = self.Wk(x).view(B, T, self.H, self.hd).transpose(1, 2) | |
| V = self.Wv(x).view(B, T, self.H, self.hd).transpose(1, 2) | |
| S = (Q @ K.transpose(-2, -1)) / (self.hd ** 0.5) + self.mask[:T, :T] | |
| A = F.softmax(S, dim=-1) | |
| O = (A @ V).transpose(1, 2).reshape(B, T, D) | |
| return self.Wo(O) | |
| class Block(nn.Module): | |
| def __init__(self, D, H, Ff, ctx): | |
| super().__init__() | |
| self.ln1 = nn.LayerNorm(D); self.attn = MHSA(D, H, ctx) | |
| self.ln2 = nn.LayerNorm(D) | |
| self.W1 = nn.Linear(D, Ff); self.W2 = nn.Linear(Ff, D) | |
| def forward(self, x): | |
| x = x + self.attn(self.ln1(x)) | |
| x = x + self.W2(F.relu(self.W1(self.ln2(x)))) | |
| return x | |
| class Transformer(nn.Module): | |
| def __init__(self, V, ctx, D, H, L, Ff): | |
| super().__init__() | |
| self.V, self.ctx, self.D = V, ctx, D | |
| self.H, self.L, self.Ff = H, L, Ff | |
| self.tok_emb = nn.Embedding(V, D) | |
| self.pos_emb = nn.Embedding(ctx, D) | |
| self.blocks = nn.ModuleList([Block(D, H, Ff, ctx) for _ in range(L)]) | |
| self.ln_f = nn.LayerNorm(D) | |
| self.Wout = nn.Linear(D, V, bias=True) | |
| def forward(self, X): | |
| B, T = X.shape | |
| pos = torch.arange(T, device=X.device).unsqueeze(0) | |
| x = self.tok_emb(X) + self.pos_emb(pos) | |
| for b in self.blocks: x = b(x) | |
| return self.Wout(self.ln_f(x)) | |
| def focal_loss(logits, Y, M, gamma, log_V): | |
| B, T, V = logits.shape | |
| lf = logits.reshape(B*T, V) | |
| Yf = Y.reshape(B*T); Mf = M.reshape(B*T) | |
| log_probs = F.log_softmax(lf, dim=-1) | |
| probs = log_probs.exp() | |
| idx = torch.arange(B*T, device=logits.device) | |
| p_t = probs[idx, Yf].clamp(1e-9, 1.0 - 1e-9) | |
| log_pt = -torch.log(p_t) | |
| focal_val = ((1.0 - p_t) ** gamma) * log_pt | |
| Msum = Mf.sum() + 1e-9 | |
| loss = (focal_val * Mf).sum() / Msum | |
| loss_norm = loss / log_V | |
| # grad: s * (probs - onehot) * mask / Msum | |
| s = ((1.0 - p_t) ** gamma) + gamma * p_t * ((1.0 - p_t) ** (gamma - 1.0)) * log_pt | |
| s = torch.nan_to_num(s, nan=0.0, posinf=0.0, neginf=0.0) | |
| grad_logits = probs.clone() | |
| grad_logits[idx, Yf] -= 1.0 | |
| grad_logits = grad_logits * Mf.unsqueeze(1) * s.unsqueeze(1) / Msum | |
| return loss_norm, grad_logits.reshape(B, T, V) | |
| def to_b64(t): | |
| return base64.b64encode(t.detach().cpu().numpy().astype("<f4").tobytes()).decode("ascii") | |
| def export_weights(model, vocab): | |
| tokens_por_id = [None] * len(vocab) | |
| for tok, idx in vocab.items(): tokens_por_id[idx] = tok | |
| blocks_data = [] | |
| for b in model.blocks: | |
| blocks_data.append({ | |
| "ln1_gamma": to_b64(b.ln1.weight), "ln1_beta": to_b64(b.ln1.bias), | |
| "ln2_gamma": to_b64(b.ln2.weight), "ln2_beta": to_b64(b.ln2.bias), | |
| "attn": { | |
| "Wq": to_b64(b.attn.Wq.weight.T), "Wk": to_b64(b.attn.Wk.weight.T), | |
| "Wv": to_b64(b.attn.Wv.weight.T), "Wo": to_b64(b.attn.Wo.weight.T), | |
| "bq": to_b64(b.attn.Wq.bias), "bk": to_b64(b.attn.Wk.bias), | |
| "bv": to_b64(b.attn.Wv.bias), "bo": to_b64(b.attn.Wo.bias), | |
| }, | |
| "W1": to_b64(b.W1.weight.T), "b1": to_b64(b.W1.bias), | |
| "W2": to_b64(b.W2.weight.T), "b2": to_b64(b.W2.bias), | |
| }) | |
| return { | |
| "config": { | |
| "arch": "transformer", "vocab_size": model.V, | |
| "context_length": model.ctx, "embed_dim": model.D, | |
| "num_heads": model.H, "num_layers": model.L, | |
| "ffn_hidden": model.Ff, "activation": "relu", | |
| }, | |
| "generation": {"temperature": 0.6, "top_k": 5, "repetition_penalty": 1.2}, | |
| "vocab": tokens_por_id, | |
| "tok_emb": to_b64(model.tok_emb.weight), | |
| "pos_emb": to_b64(model.pos_emb.weight), | |
| "ln_f_gamma": to_b64(model.ln_f.weight), | |
| "ln_f_beta": to_b64(model.ln_f.bias), | |
| "blocks": blocks_data, | |
| "output_layer": {"W": to_b64(model.Wout.weight), "b": to_b64(model.Wout.bias)}, | |
| } | |
| def train_fn(max_epochs, target_loss, progress=gr.Progress()): | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| with open(CONFIG["corpus_file"], "r", encoding="utf-8") as f: | |
| text = f.read() | |
| lines = [l.strip() for l in text.splitlines() if l.strip()] | |
| all_toks = [] | |
| for line in lines: all_toks.extend(tokenize(line)) | |
| vocab = build_vocab(all_toks, max_vocab=CONFIG["max_vocab"]) | |
| ctx = CONFIG["context_length"] | |
| X, Y, M = build_dataset(lines, vocab, ctx) | |
| V = len(vocab); log_V = float(np.log(V)) | |
| model = Transformer(V, ctx, CONFIG["embed_dim"], CONFIG["num_heads"], | |
| CONFIG["num_layers"], CONFIG["ffn_hidden"]).to(device) | |
| opt = torch.optim.Adam(model.parameters(), lr=CONFIG["learning_rate"], | |
| betas=(0.9, 0.999), eps=1e-8) | |
| n = len(X); bs = CONFIG["batch_size"] | |
| n_batches = (n + bs - 1) // bs | |
| X, Y, M = X.to(device), Y.to(device), M.to(device) | |
| t0 = time.time() | |
| log_lines = [] | |
| for ep in range(1, max_epochs + 1): | |
| perm = torch.randperm(n, device=device) | |
| Xs, Ys, Ms = X[perm], Y[perm], M[perm] | |
| ep_loss, nb = 0.0, 0 | |
| for i in range(0, n, bs): | |
| Xb, Yb, Mb = Xs[i:i+bs], Ys[i:i+bs], Ms[i:i+bs] | |
| opt.zero_grad() | |
| logits = model(Xb) | |
| loss, grad_logits = focal_loss(logits, Yb, Mb, CONFIG["focal_gamma"], log_V) | |
| logits.backward(grad_logits) | |
| opt.step() | |
| ep_loss += loss.item(); nb += 1 | |
| avg = ep_loss / nb | |
| elapsed = time.time() - t0 | |
| log_lines.append(f"Ep {ep:4d} | loss {avg:.4f} | {elapsed:6.1f}s | device={device.type}") | |
| progress(ep / max_epochs, desc=f"Ep {ep} | loss {avg:.4f}") | |
| if avg <= target_loss: break | |
| ckpt = export_weights(model, vocab) | |
| out_path = "/tmp/AI_Weights.json" | |
| with open(out_path, "w", encoding="utf-8") as f: | |
| json.dump(ckpt, f) | |
| return "\n".join(log_lines[-30:]), out_path | |
| def launch_train(max_epochs, target_loss): | |
| log, path = train_fn(int(max_epochs), float(target_loss)) | |
| return log, path | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Transformer Trainer (ZeroGPU)") | |
| with gr.Row(): | |
| ep_in = gr.Number(value=600, label="max_epochs", precision=0) | |
| tl_in = gr.Number(value=0.10, label="target_loss") | |
| btn = gr.Button("Treinar", variant="primary") | |
| log_out = gr.Textbox(label="Log", lines=20) | |
| file_out = gr.File(label="AI_Weights.json") | |
| btn.click(launch_train, inputs=[ep_in, tl_in], outputs=[log_out, file_out]) | |
| demo.launch() |