""" Main Training Script for 5M Terminal & Multilingual Language Model Runs on Kaggle Dual GPUs or local hardware. Automatically uploads model to HuggingFace Hub. """ import os import math import time import random import json import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from transformers import PreTrainedTokenizerFast from huggingface_hub import HfApi, login # ========================================== # 1. Environment & Credentials Configuration # ========================================== HF_TOKEN = os.environ.get("HF_TOKEN", "YOUR_HF_TOKEN") HF_REPO_ID = os.environ.get("HF_REPO_ID", "kipasyangin5/terminal-lang-5m") if HF_TOKEN: try: login(token=HF_TOKEN) print(f"[HF Login] Authenticated successfully as '{HF_REPO_ID.split('/')[0]}'.") except Exception as e: print(f"[HF Login Warning] Could not login: {e}") # ========================================== # 2. Model Architecture (5.0M Parameters) # ========================================== class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight class SwiGLUMLP(nn.Module): def __init__(self, dim, inter_dim): super().__init__() self.gate_proj = nn.Linear(dim, inter_dim, bias=False) self.up_proj = nn.Linear(dim, inter_dim, bias=False) self.down_proj = nn.Linear(inter_dim, dim, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class CausalSelfAttention(nn.Module): def __init__(self, dim, n_head, max_seq_len=512): super().__init__() self.dim = dim self.n_head = n_head self.head_dim = dim // n_head self.q_proj = nn.Linear(dim, dim, bias=False) self.k_proj = nn.Linear(dim, dim, bias=False) self.v_proj = nn.Linear(dim, dim, bias=False) self.out_proj = nn.Linear(dim, dim, bias=False) def forward(self, x): B, T, C = x.shape q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2) # PyTorch Scaled Dot Product Attention with Causal Mask y = F.scaled_dot_product_attention(q, k, v, is_causal=True) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.out_proj(y) class TransformerBlock(nn.Module): def __init__(self, dim, n_head, inter_dim, max_seq_len=512): super().__init__() self.attn = CausalSelfAttention(dim, n_head, max_seq_len) self.mlp = SwiGLUMLP(dim, inter_dim) self.norm1 = RMSNorm(dim) self.norm2 = RMSNorm(dim) def forward(self, x): x = x + self.attn(self.norm1(x)) x = x + self.mlp(self.norm2(x)) return x class TerminalLM5M(nn.Module): def __init__(self, vocab_size=4096, dim=256, n_layer=6, n_head=8, inter_dim=512, max_seq_len=512): super().__init__() self.vocab_size = vocab_size self.dim = dim self.max_seq_len = max_seq_len self.tok_embeddings = nn.Embedding(vocab_size, dim) self.pos_embeddings = nn.Embedding(max_seq_len, dim) self.layers = nn.ModuleList([ TransformerBlock(dim, n_head, inter_dim, max_seq_len) for _ in range(n_layer) ]) self.norm = RMSNorm(dim) self.lm_head = nn.Linear(dim, vocab_size, bias=False) # Weight Tying for memory efficiency & parameter budget self.lm_head.weight = self.tok_embeddings.weight # Parameter Initialization self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, input_ids, targets=None): B, T = input_ids.shape device = input_ids.device pos = torch.arange(0, T, dtype=torch.long, device=device) h = self.tok_embeddings(input_ids) + self.pos_embeddings(pos) for layer in self.layers: h = layer(h) h = self.norm(h) logits = self.lm_head(h) loss = None if targets is not None: loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1)) return logits, loss # ========================================== # 3. Dataset & Data Loader Construction # ========================================== from terminal_dataset import COMMAND_TEMPLATES, SHELL_INTERACTIONS def generate_mixed_corpus(num_terminal=20000, num_english=10000, num_multilingual=20000): corpus = [] # 1. Terminal Commands for _ in range(num_terminal): qa = random.choice(COMMAND_TEMPLATES) fmt = random.choice([ f"User: {qa[0]}\nAssistant: Run `{qa[1]}`\n", f"Question: {qa[0]}\nAnswer:\n```bash\n{qa[1]}\n```\n", f"$ {qa[1]}\n# Output: success\n" ]) corpus.append(fmt) # 2. Shell session interactions for _ in range(num_terminal // 2): s = random.choice(SHELL_INTERACTIONS) corpus.append(f"```session\n{s}\n```\n") # 3. English General Language (30%) en_samples = [ "The Linux kernel provides low-level hardware abstraction, process management, and networking capabilities.", "Version control systems like Git allow multiple developers to collaborate on codebases seamlessly.", "Computer networks transmit data packets across interconnected routers using TCP and IP protocols.", "Machine learning models optimize parameters using loss gradients computed via automatic differentiation.", "Shell scripts automate repetitive terminal tasks using conditional loops and system environment variables.", "Cloud infrastructure scales computational workloads across distributed server clusters efficiently." ] for _ in range(num_english): corpus.append(random.choice(en_samples) + "\n") # 4. Multilingual General Language (70% Non-English) multi_samples = [ # Indonesian "Model bahasa ini dilatih untuk mengenali perintah baris terminal Linux dan bahasa umum secara efisien.", "Perintah cd digunakan untuk berpindah direktori, sedangkan ls -la menampilkan semua berkas tersembunyi.", "Pengembangan sistem operasi berbasis Linux memungkinkan fleksibilitas tinggi bagi pengembang perangkat lunak.", # Spanish "El comando cd permite cambiar de directorio y ls -la muestra todos los archivos ocultos en la carpeta.", "Los modelos de lenguaje pequeños pueden ejecutarse eficientemente en dispositivos locales y servidores.", # French "La commande cd permet de changer de répertoire et ls -la affiche tous les fichiers cachés.", "Les modèles informatiques modernes permettent d'automatiser le traitement du langage naturel.", # German "Der Befehl cd wechselt das Verzeichnis und ls -la zeigt alle versteckten Dateien an.", "Künstliche Intelligenz optimiert die Verarbeitung von Befehlen auf modernen Betriebssystemen." ] for _ in range(num_multilingual): corpus.append(random.choice(multi_samples) + "\n") random.shuffle(corpus) return corpus class TextDataset(Dataset): def __init__(self, token_ids, seq_len=256): self.seq_len = seq_len # Pack tokens into fixed length chunks self.num_samples = (len(token_ids) - 1) // seq_len self.inputs = [] self.targets = [] for i in range(self.num_samples): start = i * seq_len end = start + seq_len self.inputs.append(token_ids[start:end]) self.targets.append(token_ids[start+1:end+1]) def __len__(self): return len(self.inputs) def __getitem__(self, idx): return torch.tensor(self.inputs[idx], dtype=torch.long), torch.tensor(self.targets[idx], dtype=torch.long) # ========================================== # 4. Main Training Routine # ========================================== def train( max_steps=2000, batch_size=32, seq_len=256, lr=1e-3, save_hf=True ): device = "cuda" if torch.cuda.is_available() else "cpu" print(f"=== Starting Training for 5M Terminal LM on {device.upper()} ===") # Load Tokenizer tok_dir = "tokenizer_5m" if not os.path.exists(tok_dir): from tokenizer_builder import build_tokenizer build_tokenizer(tok_dir) tokenizer = PreTrainedTokenizerFast.from_pretrained(tok_dir) vocab_size = len(tokenizer) print(f"[Dataset] Tokenizer loaded with vocab_size = {vocab_size}") # Build Mixed Corpus & Tokenize print("[Dataset] Building training corpus...") corpus = generate_mixed_corpus() full_text = "\n".join(corpus) print(f"[Dataset] Full text character length: {len(full_text):,}") tokens = tokenizer.encode(full_text) print(f"[Dataset] Total encoded tokens: {len(tokens):,}") dataset = TextDataset(tokens, seq_len=seq_len) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True) print(f"[Dataset] Total training batches per epoch: {len(dataloader)}") # Instantiate Model model = TerminalLM5M( vocab_size=vocab_size, dim=256, n_layer=6, n_head=8, inter_dim=512, max_seq_len=seq_len ).to(device) param_count = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"[Model] Total Trainable Parameters: {param_count:,} (~{param_count/1e6:.2f}M)") # Optimizer & Scheduler try: optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01, fused=True) except Exception: optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01) scaler = torch.amp.GradScaler('cuda') if device == "cuda" else None # Training Loop model.train() step = 0 t0 = time.time() data_iter = iter(dataloader) while step < max_steps: try: x, y = next(data_iter) except StopIteration: data_iter = iter(dataloader) x, y = next(data_iter) x, y = x.to(device), y.to(device) optimizer.zero_grad() if device == "cuda": with torch.amp.autocast('cuda', dtype=torch.float16): logits, loss = model(x, y) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() else: logits, loss = model(x, y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() step += 1 if step % 50 == 0 or step == 1: t1 = time.time() dt = t1 - t0 t0 = t1 tokens_per_sec = (50 * batch_size * seq_len) / (dt if dt > 0 else 1.0) ppl = math.exp(min(loss.item(), 20.0)) print(f"Step {step:4d}/{max_steps} | Loss: {loss.item():.4f} | PPL: {ppl:.2f} | Speed: {tokens_per_sec:.0f} tok/s") print("\n=== Training Completed Successfully ===") # Save Model & Tokenizer locally output_dir = "saved_5m_model" os.makedirs(output_dir, exist_ok=True) torch.save(model.state_dict(), os.path.join(output_dir, "model.pt")) tokenizer.save_pretrained(output_dir) print(f"[Save] Model and Tokenizer saved to '{output_dir}'.") # Upload to HuggingFace Hub if configured if save_hf and HF_TOKEN: try: print(f"[HuggingFace] Uploading model to repository '{HF_REPO_ID}'...") api = HfApi() api.create_repo(repo_id=HF_REPO_ID, exist_ok=True) api.upload_folder( folder_path=output_dir, repo_id=HF_REPO_ID, commit_message=f"Upload trained 5M Terminal LM (Steps: {max_steps}, Loss: {loss.item():.4f})" ) print(f"🚀 [SUCCESS] Model successfully uploaded to https://huggingface.co/{HF_REPO_ID}") except Exception as e: print(f"[HuggingFace Upload Warning] Could not upload to HF: {e}") if __name__ == "__main__": train(max_steps=100 if not torch.cuda.is_available() else 3000)