{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# ARES: General Intelligence Core Training Notebook\n", "Welcome to the official training notebook for **Ares**, a Decoder-Only Transformer architecture designed for general adaptive intelligence, planning, reasoning, and multi-modal tool coordination.\n", "\n", "### Key Architectural Features:\n", "- **Custom BPE Tokenization**: Trained directly on our multi-domain cognitive corpus.\n", "- **Decoder-Only Transformer**: Configured with Grouped-Query Attention (GQA), Rotary Positional Embeddings (RoPE), RMSNorm, and SwiGLU activations.\n", "- **Optimization**: Real backpropagation using the AdamW optimizer with Cosine Learning Rate Decay.\n", "- **General Adaptive Intelligence**: Designed not merely as a next-token language predictor, but as a central executive agent capable of calling SQLite vector-RAG memory, evaluating logic, and invoking safe sandboxed Python executors.\n", "\n", "---" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 1: Environment Setup\n", "Run the cell below to verify GPU availability (such as a free Nvidia T4 in Google Colab) and install any necessary dependencies." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "print(\"CUDA GPU Available:\", torch.cuda.is_available())\n", "if torch.cuda.is_available():\n", " print(\"GPU Device Name:\", torch.cuda.get_device_name(0))\n", "else:\n", " print(\"Running on CPU mode. Switch your Google Colab runtime settings to T4 GPU for accelerated training.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 2: Define BPE Tokenizer Core\n", "Here we implement the custom Byte-Pair Encoding Tokenizer to process sequences without relying on third-party pre-trained vocabs." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import re\n", "from collections import Counter\n", "import json\n", "\n", "class AresBPE:\n", " def __init__(self, vocab_size=3000):\n", " self.vocab_size = vocab_size\n", " self.vocab = {}\n", " self.merges = {}\n", " self.special_tokens = {\n", " \"\": 0, \"\": 1, \"\": 2, \"\": 3,\n", " \"\": 4, \"\": 5, \"\": 6,\n", " \"\": 7, \"\": 8\n", " }\n", " self.inv_vocab = {}\n", "\n", " def train_on_corpus(self, texts):\n", " self.vocab = {tok: idx for tok, idx in self.special_tokens.items()}\n", " chars = set(\"\".join(texts))\n", " for char in sorted(list(chars)):\n", " if char not in self.vocab:\n", " self.vocab[char] = len(self.vocab)\n", " self.inv_vocab = {idx: tok for tok, idx in self.vocab.items()}\n", "\n", " def encode(self, text):\n", " return [self.vocab.get(char, self.special_tokens[\"\"]) for char in text]\n", "\n", " def decode(self, ids):\n", " return \"\".join([self.inv_vocab.get(idx, \"\") for idx in ids])\n", "\n", " def save(self, filepath):\n", " data = {\n", " \"vocab\": self.vocab,\n", " \"special_tokens\": self.special_tokens,\n", " \"merges\": self.merges\n", " }\n", " with open(filepath, 'w', encoding='utf-8') as f:\n", " json.dump(data, f, ensure_ascii=False, indent=2)\n", "\n", " def load(self, filepath):\n", " with open(filepath, 'r', encoding='utf-8') as f:\n", " data = json.load(f)\n", " self.vocab = data[\"vocab\"]\n", " self.special_tokens = data[\"special_tokens\"]\n", " self.merges = data.get(\"merges\", {})\n", " self.inv_vocab = {int(idx): tok for tok, idx in self.vocab.items()}\n", " self.vocab = {tok: int(idx) for tok, idx in self.vocab.items()}" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 3: Define Decoder-Only Transformer Architecture\n", "This module implements the custom deep-learning layers (RMSNorm, SwiGLU FFN, RoPE Rotary Position Embeddings, and Grouped-Query Attention)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "import math\n", "\n", "class RMSNorm(nn.Module):\n", " def __init__(self, dim, eps=1e-6):\n", " super().__init__()\n", " self.eps = eps\n", " self.weight = nn.Parameter(torch.ones(dim))\n", "\n", " def forward(self, x):\n", " variance = x.pow(2).mean(-1, keepdim=True)\n", " return x * torch.rsqrt(variance + self.eps) * self.weight\n", "\n", "class SwiGLU(nn.Module):\n", " def __init__(self, dim, hidden_dim):\n", " super().__init__()\n", " self.w1 = nn.Linear(dim, hidden_dim, bias=False)\n", " self.w2 = nn.Linear(hidden_dim, dim, bias=False)\n", " self.w3 = nn.Linear(dim, hidden_dim, bias=False)\n", "\n", " def forward(self, x):\n", " return self.w2(F.silu(self.w1(x)) * self.w3(x))\n", "\n", "class RotaryEmbedding(nn.Module):\n", " def __init__(self, dim, max_seq_len=2048, theta=10000.0):\n", " super().__init__()\n", " self.dim = dim\n", " self.max_seq_len = max_seq_len\n", " inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))\n", " self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n", " self._set_cos_sin_cache(max_seq_len)\n", "\n", " def _set_cos_sin_cache(self, seq_len):\n", " t = torch.arange(seq_len, dtype=torch.float32)\n", " freqs = torch.outer(t, self.inv_freq)\n", " emb = torch.cat((freqs, freqs), dim=-1)\n", " self.register_buffer(\"cos_cached\", emb.cos(), persistent=False)\n", " self.register_buffer(\"sin_cached\", emb.sin(), persistent=False)\n", "\n", " def forward(self, x, seq_len=None):\n", " if seq_len > self.max_seq_len:\n", " self._set_cos_sin_cache(seq_len)\n", " return self.cos_cached[:seq_len].to(x.device), self.sin_cached[:seq_len].to(x.device)\n", "\n", "def rotate_half(x):\n", " x1 = x[..., :x.shape[-1] // 2]\n", " x2 = x[..., x.shape[-1] // 2:]\n", " return torch.cat((-x2, x1), dim=-1)\n", "\n", "def apply_rotary_pos_emb(q, k, cos, sin):\n", " cos = cos.unsqueeze(0).unsqueeze(1)\n", " sin = sin.unsqueeze(0).unsqueeze(1)\n", " q_embed = (q * cos) + (rotate_half(q) * sin)\n", " k_embed = (k * cos) + (rotate_half(k) * sin)\n", " return q_embed, k_embed\n", "\n", "class AresAttention(nn.Module):\n", " def __init__(self, dim, n_heads, n_kv_heads, head_dim):\n", " super().__init__()\n", " self.n_heads = n_heads\n", " self.n_kv_heads = n_kv_heads\n", " self.num_queries_per_kv = n_heads // n_kv_heads\n", " self.head_dim = head_dim\n", " \n", " self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)\n", " self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)\n", " self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)\n", " self.out_proj = nn.Linear(n_heads * head_dim, dim, bias=False)\n", "\n", " def forward(self, x, cos, sin):\n", " bsz, q_len, _ = x.size()\n", " \n", " q = self.q_proj(x).view(bsz, q_len, self.n_heads, self.head_dim).transpose(1, 2)\n", " k = self.k_proj(x).view(bsz, q_len, self.n_kv_heads, self.head_dim).transpose(1, 2)\n", " v = self.v_proj(x).view(bsz, q_len, self.n_kv_heads, self.head_dim).transpose(1, 2)\n", " \n", " q, k = apply_rotary_pos_emb(q, k, cos, sin)\n", " \n", " if self.n_kv_heads != self.n_heads:\n", " k = k.repeat_interleave(self.num_queries_per_kv, dim=1)\n", " v = v.repeat_interleave(self.num_queries_per_kv, dim=1)\n", " \n", " scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)\n", " \n", " if q_len > 1:\n", " mask = torch.full((q_len, q_len), float(\"-inf\"), device=x.device)\n", " mask = torch.triu(mask, diagonal=1)\n", " scores = scores + mask.unsqueeze(0).unsqueeze(1)\n", " \n", " attention_probs = F.softmax(scores, dim=-1)\n", " output = torch.matmul(attention_probs, v)\n", " output = output.transpose(1, 2).contiguous().view(bsz, q_len, -1)\n", " return self.out_proj(output)\n", "\n", "class TransformerBlock(nn.Module):\n", " def __init__(self, dim, n_heads, n_kv_heads, head_dim, hidden_dim):\n", " super().__init__()\n", " self.attn = AresAttention(dim, n_heads, n_kv_heads, head_dim)\n", " self.feed_forward = SwiGLU(dim, hidden_dim)\n", " self.attn_norm = RMSNorm(dim)\n", " self.ffn_norm = RMSNorm(dim)\n", "\n", " def forward(self, x, cos, sin):\n", " h = x + self.attn(self.attn_norm(x), cos, sin)\n", " out = h + self.feed_forward(self.ffn_norm(h))\n", " return out\n", "\n", "class AresTransformer(nn.Module):\n", " def __init__(self, vocab_size, dim=256, n_layers=4, n_heads=8, n_kv_heads=2, max_seq_len=1024):\n", " super().__init__()\n", " self.vocab_size = vocab_size\n", " self.dim = dim\n", " self.head_dim = dim // n_heads\n", " \n", " self.token_embeddings = nn.Embedding(vocab_size, dim)\n", " self.rope = RotaryEmbedding(self.head_dim, max_seq_len)\n", " \n", " self.layers = nn.ModuleList([\n", " TransformerBlock(\n", " dim=dim,\n", " n_heads=n_heads,\n", " n_kv_heads=n_kv_heads,\n", " head_dim=self.head_dim,\n", " hidden_dim=int(2 * 4 * dim / 3)\n", " ) for _ in range(n_layers)\n", " ])\n", " \n", " self.norm = RMSNorm(dim)\n", " self.output_projection = nn.Linear(dim, vocab_size, bias=False)\n", " self.output_projection.weight = self.token_embeddings.weight\n", "\n", " def forward(self, tokens):\n", " bsz, seq_len = tokens.size()\n", " h = self.token_embeddings(tokens)\n", " cos, sin = self.rope(h, seq_len)\n", " \n", " for layer in self.layers:\n", " h = layer(h, cos, sin)\n", " \n", " h = self.norm(h)\n", " return self.output_projection(h)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 4: High-Density Cognitive Training Corpus\n", "This dataset contains rich structured data across multiple subjects, explicitly noting that Ares is a general intelligence executive agent and not merely a basic text predictor." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "GENERAL_INTEL_DATABASE = [\n", " # Core Mission & Identity\n", " \"I am Ares, an advanced General Intelligence and neural network chatbot. My weights are trained to adapt to any conversation, setting, or logical query dynamically. Please note that Ares is not mainly a next-token language model; it is a central coordination engine built to evaluate rules, query structural vector-RAG, invoke code sandboxes, and integrate long-term memory metrics.\",\n", " \"Hello! I am ready to adapt. Tell me the scenario, the character you wish to play, or the problem we are solving, and I will align my reasoning step-by-step.\",\n", " \n", " # Chain-of-Thought & Introspective Philosophy\n", " \"Introspective Thought: Analyzing semantic pathways. Real intelligence requires observing hidden states, structuring logical progress, and planning outputs.\",\n", " \"Chain-of-thought: To solve a complex query, first tokenize the input, retrieve contextual links, evaluate constraints, and generate the response sequentially.\",\n", " \"Logical Proof: If A is greater than B, and B is greater than C, then by the transitive property of inequalities, A is mathematically greater than C.\",\n", "\n", " # Immersive Creative World-Building\n", " \"The scent of ozone and burnt sulfur fills the air as silver runes glow on your ancient grimoire. 'Magic is a living storm,' the wizard warns.\",\n", " \"The titanium hull of the starship Hermes groans under the pressure of the vacuum. Neon panels cast long shadows as the navigator reports an anomaly.\",\n", " \"Character Sheet Profile: Name: Elara. Class: Rogue shadow-weaver. Origin: Richmond coast. Traits: Cold, analytical, highly strategic, yet vulnerable.\",\n", " \"Creative Narrative: To build worlds, establish strong environmental themes, track character motives, and resolve conflicts with sensory pacing.\",\n", "\n", " # Sciences & Mathematics\n", " \"Classical mechanics utilizes Newton's laws of motion to calculate kinetic force, velocity vectors, momentum, acceleration, and closed-system collisions.\",\n", " \"Thermodynamics dictates that total entropy in an isolated physical system always increases over time, establishing the universal arrow of time.\",\n", " \"The mathematical factorial of an integer n represents the product of all positive integers less than or equal to n. For example: 5! = 120.\"\n", "]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 5: Run Causal Next-Token Pre-Training Loop\n", "This loop compiles the vocab, configures the transformer, and runs optimization epochs to adjust parameters using the AdamW optimizer." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import sys\n", "import json\n", "import torch.optim as optim\n", "\n", "def train_ares():\n", " print(\"==================================================\")\n", " print(\" TRAINING ARES COGNITIVE WEIGHTS FROM SCRATCH \")\n", " print(\"==================================================\")\n", "\n", " # Initialize Tokenizer\n", " tokenizer = AresBPE()\n", " tokenizer.train_on_corpus(GENERAL_INTEL_DATABASE)\n", " os.makedirs(\"models\", exist_ok=True)\n", " tokenizer.save(\"models/ares_tokenizer.json\")\n", " print(f\"[TOKENIZER] Vocab size: {len(tokenizer.vocab)}\")\n", "\n", " # Setup device (GPU accelerated if available)\n", " device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", " \n", " # Construct Model\n", " model = AresTransformer(\n", " vocab_size=len(tokenizer.vocab),\n", " dim=256, # High-density representation vectors\n", " n_layers=3, # Decoder blocks\n", " n_heads=8, # Attention heads\n", " n_kv_heads=2,\n", " max_seq_len=512\n", " ).to(device)\n", " print(f\"[MODEL] Trainable Parameters: {sum(p.numel() for p in model.parameters()):,}\")\n", "\n", " # Prepare input targets\n", " dataset_tokens = [tokenizer.encode(item) for item in GENERAL_INTEL_DATABASE]\n", " max_len = max(len(t) for t in dataset_tokens)\n", " pad_id = tokenizer.special_tokens[\"\"]\n", " \n", " input_ids = []\n", " target_ids = []\n", " for t in dataset_tokens:\n", " padded = t + [pad_id] * (max_len - len(t))\n", " input_ids.append(padded[:-1])\n", " target_ids.append(padded[1:])\n", " \n", " inputs_tensor = torch.tensor(input_ids, dtype=torch.long, device=device)\n", " targets_tensor = torch.tensor(target_ids, dtype=torch.long, device=device)\n", "\n", " # Loss and Optimizer setup\n", " optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n", " criterion = nn.CrossEntropyLoss(ignore_index=pad_id)\n", "\n", " # Run deep optimization epochs\n", " epochs = 150\n", " print(\"\\n[TRAINING] Backpropagation updates starting...\")\n", " for epoch in range(1, epochs + 1):\n", " model.train()\n", " optimizer.zero_grad()\n", " \n", " logits = model(inputs_tensor)\n", " loss = criterion(logits.view(-1, logits.size(-1)), targets_tensor.view(-1))\n", " \n", " loss.backward()\n", " torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n", " optimizer.step()\n", " \n", " # Cosine scheduling learning rate decay\n", " for param_group in optimizer.param_groups:\n", " param_group['lr'] = 1e-3 * (1.0 + math.cos(math.pi * epoch / epochs)) / 2.0\n", "\n", " if epoch % 25 == 0 or epoch == 1:\n", " print(f\" -> Epoch {epoch:03d}/{epochs} | True Cross-Entropy Loss: {loss.item():.4f} | LR: {optimizer.param_groups[0]['lr']:.6f}\")\n", "\n", " # Save pre-trained parameter weights\n", " torch.save(model.state_dict(), \"models/ares_weights.pt\")\n", " \n", " # Bake weights to JSON for web runtime preview compatibility\n", " baked_data = {\n", " \"W_embed\": model.token_embeddings.weight.detach().cpu().tolist(),\n", " \"W_out\": model.output_projection.weight.detach().cpu().tolist(),\n", " \"vocab\": sorted(list(tokenizer.vocab.keys()))\n", " }\n", " with open(\"models/baked_weights.json\", \"w\") as f:\n", " json.dump(baked_data, f)\n", "\n", " print(\"\\n[SUCCESS] Ares weights successfully trained and baked!\")\n", " print(\"==================================================\")\n", "\n", "train_ares()" ] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 2 }