training_code/train.py

"""
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)