{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# OFFICIAL Ares Colab Training Run\n", "\n", "This notebook is the clean, official Colab run for training Ares so it can generate responses from **its own checkpoint weights**.\n", "\n", "It is designed to be followed top-to-bottom.\n", "\n", "## What this notebook trains on\n", "\n", "1. **Official/public dataset:** `wikimedia/wikipedia`, streamed from Hugging Face.\n", "2. **Optional public language dataset:** `roneneldan/TinyStories`, streamed from Hugging Face for simpler narrative language.\n", "3. **Ares generated curriculum:** deterministic machine-learning-process examples.\n", "4. **Ares roleplay dataset:** thousands of deterministic roleplay chat scenarios.\n", "5. **Ares SFT pass:** optional roleplay-focused supervised fine-tuning after base training.\n", "\n", "## What real generation means here\n", "\n", "The Hugging Face Static Space UI cannot run server-side PyTorch. Real Ares generation happens in this notebook after training. The generation cells load your checkpoint and sample tokens from Ares weights. No external AI API is used for the model brain.\n", "\n", "## Run this carefully\n", "\n", "Start with the default `PROFILE = \"official_30m_512\"`. If you hit CUDA out-of-memory, switch to `smoke_10m`. If it works and you want more context, switch to `context_30m_1024`.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. Enable GPU\n", "\n", "In Colab, select:\n", "\n", "```text\n", "Runtime \u2192 Change runtime type \u2192 GPU\n", "```\n", "\n", "Then run this cell. If CUDA is false, stop and enable GPU before training.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os, sys, subprocess, json, time, textwrap\n", "from pathlib import Path\n", "\n", "try:\n", " import torch\n", " print('torch:', torch.__version__)\n", " print('cuda available:', torch.cuda.is_available())\n", " if torch.cuda.is_available():\n", " print('gpu:', torch.cuda.get_device_name(0))\n", " print('vram GB:', round(torch.cuda.get_device_properties(0).total_memory / 1e9, 2))\n", " else:\n", " print('WARNING: CUDA is false. Use Runtime \u2192 Change runtime type \u2192 GPU.')\n", "except Exception as e:\n", " print('Torch check before install failed:', repr(e))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. Install dependencies\n", "\n", "Colab already includes PyTorch. This installs the packages Ares needs without intentionally replacing Colab's CUDA torch build.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!pip -q install -U tokenizers datasets safetensors huggingface_hub tqdm numpy\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. Load the official Ares Space repo\n", "\n", "This clones the official Static Space repo. If you uploaded `ares-static-space.zip` manually, it will unpack that instead.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from pathlib import Path\n", "import os, sys, subprocess, zipfile, shutil\n", "\n", "PROJECT = Path('/content/ares-static-space')\n", "ZIP_PATH = Path('/content/ares-static-space.zip')\n", "REPO_URL = 'https://huggingface.co/spaces/jacmor64/ares-static-lab'\n", "\n", "if not (PROJECT / 'ares_core').exists():\n", " if ZIP_PATH.exists():\n", " print('Unpacking uploaded ares-static-space.zip...')\n", " with zipfile.ZipFile(ZIP_PATH) as z:\n", " z.extractall('/content')\n", " else:\n", " print('Cloning official Ares Space repo...')\n", " subprocess.run(['git', 'clone', '--depth', '1', REPO_URL, str(PROJECT)], check=True)\n", "\n", "assert (PROJECT / 'ares_core').exists(), 'Ares project files were not found.'\n", "os.chdir(PROJECT)\n", "sys.path.insert(0, str(PROJECT))\n", "print('PROJECT =', PROJECT)\n", "print('Ares core files:', len(list((PROJECT / 'ares_core').glob('*.py'))))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4. Mount Google Drive for checkpoints\n", "\n", "Training artifacts should go to Drive, not the public Static Space repo. This protects checkpoints from Colab runtime resets.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "USE_DRIVE = True\n", "\n", "if USE_DRIVE:\n", " from google.colab import drive\n", " drive.mount('/content/drive')\n", " ARTIFACTS = Path('/content/drive/MyDrive/ares_official_training')\n", "else:\n", " ARTIFACTS = PROJECT / 'artifacts'\n", "ARTIFACTS.mkdir(parents=True, exist_ok=True)\n", "print('ARTIFACTS =', ARTIFACTS)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 5. Choose training profile\n", "\n", "Use one of these:\n", "\n", "- `smoke_10m` \u2014 quickest test, smaller model and less context.\n", "- `official_30m_512` \u2014 recommended first real run.\n", "- `context_30m_1024` \u2014 bigger context; use after the 512 run works.\n", "\n", "The 30M model is still small compared to frontier LLMs, but it is the first serious Ares brain target for Colab.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "PROFILE = 'official_30m_512' # 'smoke_10m', 'official_30m_512', or 'context_30m_1024'\n", "\n", "if PROFILE == 'smoke_10m':\n", " RUN_NAME = 'ares-smoke-10m'\n", " CONFIG_PATH = PROJECT / 'configs' / 'ares_8m.json'\n", " TOKENIZER_VOCAB_SIZE = 4096\n", " WIKIPEDIA_RECORDS = 1000\n", " ML_PROCESS_RECORDS = 500\n", " ROLEPLAY_RECORDS = 2000\n", " EXTRA_HF_RECORDS = 500\n", " TRAIN_STEPS = 300\n", " BATCH_SIZE = 4\n", " GRAD_ACCUM = 1\n", " RUN_ROLEPLAY_SFT = False\n", "elif PROFILE == 'official_30m_512':\n", " RUN_NAME = 'ares-official-30m-ctx512'\n", " CONFIG_PATH = PROJECT / 'configs' / 'ares_30m.json'\n", " TOKENIZER_VOCAB_SIZE = 8192\n", " WIKIPEDIA_RECORDS = 8000\n", " ML_PROCESS_RECORDS = 4000\n", " ROLEPLAY_RECORDS = 12000\n", " EXTRA_HF_RECORDS = 2000\n", " TRAIN_STEPS = 1500\n", " BATCH_SIZE = 2\n", " GRAD_ACCUM = 4\n", " RUN_ROLEPLAY_SFT = True\n", "elif PROFILE == 'context_30m_1024':\n", " RUN_NAME = 'ares-official-30m-ctx1024'\n", " CONFIG_PATH = PROJECT / 'configs' / 'ares_30m_ctx1024.json'\n", " TOKENIZER_VOCAB_SIZE = 8192\n", " WIKIPEDIA_RECORDS = 8000\n", " ML_PROCESS_RECORDS = 4000\n", " ROLEPLAY_RECORDS = 12000\n", " EXTRA_HF_RECORDS = 2000\n", " TRAIN_STEPS = 1500\n", " BATCH_SIZE = 1\n", " GRAD_ACCUM = 8\n", " RUN_ROLEPLAY_SFT = True\n", "else:\n", " raise ValueError(f'Unknown PROFILE: {PROFILE}')\n", "\n", "# Official/public datasets.\n", "WIKIPEDIA_DATASET = 'wikimedia/wikipedia'\n", "USE_FULL_EN_WIKIPEDIA = False # set True later for full English Wikipedia; start simple first\n", "WIKIPEDIA_CONFIG = '20231101.en' if USE_FULL_EN_WIKIPEDIA else '20231101.simple'\n", "\n", "# Optional extra public HF dataset for simpler language/narrative structure.\n", "INCLUDE_TINYSTORIES = True\n", "EXTRA_HF_DATASET = 'roneneldan/TinyStories' if INCLUDE_TINYSTORIES else None\n", "EXTRA_HF_CONFIG = None\n", "EXTRA_HF_SPLIT = 'train'\n", "EXTRA_HF_TEXT_FIELD = 'text'\n", "\n", "# Optimization.\n", "LEARNING_RATE = 3e-4\n", "MIN_LR = 3e-5\n", "WARMUP_STEPS = 100\n", "EVAL_EVERY = 100\n", "EVAL_BATCHES = 30\n", "SAVE_EVERY = 250\n", "LOG_EVERY = 25\n", "DEVICE = 'auto'\n", "\n", "# SFT refinement.\n", "ROLEPLAY_SFT_STEPS = 400\n", "ROLEPLAY_SFT_LR = 1e-4\n", "\n", "print(json.dumps({\n", " 'PROFILE': PROFILE,\n", " 'RUN_NAME': RUN_NAME,\n", " 'CONFIG_PATH': str(CONFIG_PATH),\n", " 'WIKIPEDIA': f'{WIKIPEDIA_DATASET}:{WIKIPEDIA_CONFIG}',\n", " 'WIKIPEDIA_RECORDS': WIKIPEDIA_RECORDS,\n", " 'EXTRA_HF_DATASET': EXTRA_HF_DATASET,\n", " 'EXTRA_HF_RECORDS': EXTRA_HF_RECORDS,\n", " 'ML_PROCESS_RECORDS': ML_PROCESS_RECORDS,\n", " 'ROLEPLAY_RECORDS': ROLEPLAY_RECORDS,\n", " 'TOKENIZER_VOCAB_SIZE': TOKENIZER_VOCAB_SIZE,\n", " 'TRAIN_STEPS': TRAIN_STEPS,\n", " 'BATCH_SIZE': BATCH_SIZE,\n", " 'GRAD_ACCUM': GRAD_ACCUM,\n", " 'RUN_ROLEPLAY_SFT': RUN_ROLEPLAY_SFT,\n", "}, indent=2))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6. Create artifact paths\n", "\n", "Everything produced by training goes under your Drive artifact folder.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "DATA_DIR = ARTIFACTS / 'data' / RUN_NAME\n", "ROLEPLAY_JSONL = DATA_DIR / 'roleplay_dataset.jsonl'\n", "ROLEPLAY_TEXT = DATA_DIR / 'roleplay_sft.txt'\n", "ROLEPLAY_TRAIN = DATA_DIR / 'roleplay_sft_train.txt'\n", "ROLEPLAY_VAL = DATA_DIR / 'roleplay_sft_val.txt'\n", "CORPUS_PATH = DATA_DIR / 'corpus_full_wiki_ml_roleplay.txt'\n", "TRAIN_PATH = DATA_DIR / 'corpus_train.txt'\n", "VAL_PATH = DATA_DIR / 'corpus_val.txt'\n", "TOKENIZER_PATH = ARTIFACTS / 'tokenizers' / f'{RUN_NAME}_tokenizer.json'\n", "OUT_DIR = ARTIFACTS / 'checkpoints' / RUN_NAME\n", "SFT_OUT_DIR = ARTIFACTS / 'checkpoints' / f'{RUN_NAME}-roleplay-sft'\n", "RAG_DB = ARTIFACTS / 'rag' / f'{RUN_NAME}.sqlite'\n", "REPORT_PATH = OUT_DIR / 'training_report.html'\n", "SFT_REPORT_PATH = SFT_OUT_DIR / 'training_report.html'\n", "SUMMARY_PATH = ARTIFACTS / 'summaries' / f'{RUN_NAME}_summary.json'\n", "\n", "for p in [DATA_DIR, TOKENIZER_PATH.parent, OUT_DIR, SFT_OUT_DIR, RAG_DB.parent, SUMMARY_PATH.parent]:\n", " p.mkdir(parents=True, exist_ok=True)\n", "\n", "print('DATA_DIR =', DATA_DIR)\n", "print('OUT_DIR =', OUT_DIR)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 7. Generate thousands of roleplay scenarios\n", "\n", "This creates Ares' own roleplay dataset in JSONL chat format. It is deterministic, local, and does not call another AI model.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cmd = [\n", " sys.executable, '-m', 'ares_core.synthetic_curriculum', 'roleplay',\n", " '--output', str(ROLEPLAY_JSONL),\n", " '--count', str(ROLEPLAY_RECORDS),\n", " '--seed', '2026',\n", " '--jsonl',\n", "]\n", "print(' '.join(map(str, cmd)))\n", "subprocess.run(cmd, check=True)\n", "print('Roleplay JSONL MB:', round(ROLEPLAY_JSONL.stat().st_size / 1e6, 2))\n", "print('First two roleplay records:')\n", "for i, line in enumerate(ROLEPLAY_JSONL.open('r', encoding='utf-8')):\n", " print(line[:700])\n", " if i >= 1:\n", " break\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 8. Convert roleplay JSONL to Ares SFT text\n", "\n", "This creates the roleplay supervised fine-tuning file used after base training.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cmd = [\n", " sys.executable, '-m', 'ares_core.sft',\n", " '--input', str(ROLEPLAY_JSONL),\n", " '--output', str(ROLEPLAY_TEXT),\n", "]\n", "print(' '.join(map(str, cmd)))\n", "subprocess.run(cmd, check=True)\n", "print('Roleplay SFT text MB:', round(ROLEPLAY_TEXT.stat().st_size / 1e6, 2))\n", "print(ROLEPLAY_TEXT.read_text(encoding='utf-8')[:1000])\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 9. Build the full corpus: Wikipedia + TinyStories + ML + roleplay\n", "\n", "This is the main Ares base-training corpus.\n", "\n", "Default official mix:\n", "\n", "- Wikipedia records: 8,000\n", "- TinyStories records: 2,000\n", "- ML-process records: 4,000\n", "- Roleplay records: 12,000\n", "\n", "You can raise counts later after the loop succeeds.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cmd = [\n", " sys.executable, '-m', 'ares_core.mixture_build',\n", " '--output', str(CORPUS_PATH),\n", " '--local', str(PROJECT / 'data' / 'sample_corpus.txt'),\n", " '--wikipedia-dataset', WIKIPEDIA_DATASET,\n", " '--wikipedia-config', WIKIPEDIA_CONFIG,\n", " '--wikipedia-records', str(WIKIPEDIA_RECORDS),\n", " '--ml-records', str(ML_PROCESS_RECORDS),\n", " '--roleplay-records', str(ROLEPLAY_RECORDS),\n", " '--seed', '2026',\n", " '--min-chars', '60',\n", " '--max-chars', '24000',\n", "]\n", "if EXTRA_HF_DATASET and EXTRA_HF_RECORDS > 0:\n", " cmd += [\n", " '--extra-hf-dataset', EXTRA_HF_DATASET,\n", " '--extra-hf-split', EXTRA_HF_SPLIT,\n", " '--extra-hf-text-field', EXTRA_HF_TEXT_FIELD,\n", " '--extra-hf-records', str(EXTRA_HF_RECORDS),\n", " ]\n", " if EXTRA_HF_CONFIG:\n", " cmd += ['--extra-hf-config', EXTRA_HF_CONFIG]\n", "print(' '.join(map(str, cmd)))\n", "subprocess.run(cmd, check=True)\n", "print('Corpus MB:', round(CORPUS_PATH.stat().st_size / 1e6, 2))\n", "manifest = Path(str(CORPUS_PATH) + '.manifest.json')\n", "print('Manifest:')\n", "print(manifest.read_text()[:2000])\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 10. Split train/validation for base corpus\n", "\n", "This is required. Do not train without validation.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cmd = [\n", " sys.executable, '-m', 'ares_core.split_corpus',\n", " '--input', str(CORPUS_PATH),\n", " '--train-output', str(TRAIN_PATH),\n", " '--val-output', str(VAL_PATH),\n", " '--val-ratio', '0.02',\n", " '--min-val-records', '100',\n", "]\n", "print(' '.join(map(str, cmd)))\n", "subprocess.run(cmd, check=True)\n", "print('Train MB:', round(TRAIN_PATH.stat().st_size / 1e6, 2))\n", "print('Val MB:', round(VAL_PATH.stat().st_size / 1e6, 2))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 11. Split roleplay SFT data\n", "\n", "This keeps a held-out roleplay validation set for the SFT phase.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cmd = [\n", " sys.executable, '-m', 'ares_core.split_corpus',\n", " '--input', str(ROLEPLAY_TEXT),\n", " '--train-output', str(ROLEPLAY_TRAIN),\n", " '--val-output', str(ROLEPLAY_VAL),\n", " '--val-ratio', '0.02',\n", " '--min-val-records', '100',\n", "]\n", "print(' '.join(map(str, cmd)))\n", "subprocess.run(cmd, check=True)\n", "print('Roleplay train MB:', round(ROLEPLAY_TRAIN.stat().st_size / 1e6, 2))\n", "print('Roleplay val MB:', round(ROLEPLAY_VAL.stat().st_size / 1e6, 2))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 12. Train the BPE tokenizer from scratch\n", "\n", "Ares' tokenizer is trained on the actual corpus. No pretrained tokenizer is used.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cmd = [\n", " sys.executable, '-m', 'ares_core.tokenizer_train',\n", " '--input', str(TRAIN_PATH), str(ROLEPLAY_TRAIN),\n", " '--output', str(TOKENIZER_PATH),\n", " '--vocab-size', str(TOKENIZER_VOCAB_SIZE),\n", " '--min-frequency', '2',\n", "]\n", "print(' '.join(map(str, cmd)))\n", "subprocess.run(cmd, check=True)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 13. Estimate tokens and training size\n", "\n", "This cell tells you approximately how much tokenized data you have and how many tokens each optimizer step sees.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from tokenizers import Tokenizer\n", "from ares_core.config import AresConfig, estimate_parameters\n", "\n", "cfg = AresConfig.from_json(CONFIG_PATH)\n", "tok = Tokenizer.from_file(str(TOKENIZER_PATH))\n", "\n", "def count_tokens(path):\n", " text = Path(path).read_text(encoding='utf-8', errors='ignore')\n", " return len(tok.encode(text).ids)\n", "\n", "train_tokens = count_tokens(TRAIN_PATH)\n", "val_tokens = count_tokens(VAL_PATH)\n", "roleplay_tokens = count_tokens(ROLEPLAY_TRAIN)\n", "params = estimate_parameters(cfg)\n", "effective_tokens_per_step = cfg.max_seq_len * BATCH_SIZE * GRAD_ACCUM\n", "planned_tokens = effective_tokens_per_step * TRAIN_STEPS\n", "\n", "print(json.dumps({\n", " 'model_name': cfg.model_name,\n", " 'params_estimate': params,\n", " 'context_window': cfg.max_seq_len,\n", " 'tokenizer_vocab': tok.get_vocab_size(),\n", " 'train_tokens': train_tokens,\n", " 'val_tokens': val_tokens,\n", " 'roleplay_sft_train_tokens': roleplay_tokens,\n", " 'effective_tokens_per_step': effective_tokens_per_step,\n", " 'planned_base_training_tokens_seen': planned_tokens,\n", "}, indent=2))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 14. Verify model size\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "subprocess.run([sys.executable, 'scripts/estimate_params.py'], check=True)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 15. Train Ares base brain\n", "\n", "This is the main base training loop.\n", "\n", "Watch these values:\n", "\n", "- `loss`\n", "- `loss_ema`\n", "- `val_loss`\n", "- `val_perplexity`\n", "- `saved_best`\n", "\n", "The best checkpoint is saved as `ckpt_best.pt`.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cmd = [\n", " sys.executable, '-m', 'ares_core.train',\n", " '--config', str(CONFIG_PATH),\n", " '--tokenizer', str(TOKENIZER_PATH),\n", " '--train', str(TRAIN_PATH),\n", " '--val', str(VAL_PATH),\n", " '--out', str(OUT_DIR),\n", " '--steps', str(TRAIN_STEPS),\n", " '--batch-size', str(BATCH_SIZE),\n", " '--grad-accum', str(GRAD_ACCUM),\n", " '--lr', str(LEARNING_RATE),\n", " '--min-lr', str(MIN_LR),\n", " '--warmup-steps', str(WARMUP_STEPS),\n", " '--eval-every', str(EVAL_EVERY),\n", " '--eval-batches', str(EVAL_BATCHES),\n", " '--save-every', str(SAVE_EVERY),\n", " '--log-every', str(LOG_EVERY),\n", " '--device', DEVICE,\n", "]\n", "print(' '.join(map(str, cmd)))\n", "subprocess.run(cmd, check=True)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 16. Generate base training report\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cmd = [\n", " sys.executable, '-m', 'ares_core.train_report',\n", " '--log', str(OUT_DIR / 'train_log.jsonl'),\n", " '--out', str(REPORT_PATH),\n", " '--title', f'Ares base training report: {RUN_NAME}',\n", "]\n", "subprocess.run(cmd, check=True)\n", "print('Report:', REPORT_PATH)\n", "try:\n", " from IPython.display import IFrame, display\n", " display(IFrame(str(REPORT_PATH), width='100%', height=720))\n", "except Exception as e:\n", " print('Inline report display skipped:', e)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 17. Optional roleplay SFT\n", "\n", "This resumes from the best base checkpoint and refines Ares on roleplay chat data. It resets optimizer state so this becomes a clean post-training phase.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "BASE_CKPT = OUT_DIR / 'ckpt_best.pt'\n", "if not BASE_CKPT.exists():\n", " BASE_CKPT = OUT_DIR / 'ckpt_last.pt'\n", "print('BASE_CKPT =', BASE_CKPT)\n", "\n", "if RUN_ROLEPLAY_SFT:\n", " cmd = [\n", " sys.executable, '-m', 'ares_core.train',\n", " '--config', str(CONFIG_PATH),\n", " '--tokenizer', str(TOKENIZER_PATH),\n", " '--train', str(ROLEPLAY_TRAIN),\n", " '--val', str(ROLEPLAY_VAL),\n", " '--out', str(SFT_OUT_DIR),\n", " '--resume', str(BASE_CKPT),\n", " '--resume-model-only',\n", " '--reset-step-on-resume',\n", " '--steps', str(ROLEPLAY_SFT_STEPS),\n", " '--batch-size', str(BATCH_SIZE),\n", " '--grad-accum', str(GRAD_ACCUM),\n", " '--lr', str(ROLEPLAY_SFT_LR),\n", " '--min-lr', str(ROLEPLAY_SFT_LR / 10),\n", " '--warmup-steps', '50',\n", " '--eval-every', str(max(50, ROLEPLAY_SFT_STEPS // 4)),\n", " '--eval-batches', str(EVAL_BATCHES),\n", " '--save-every', str(max(100, ROLEPLAY_SFT_STEPS // 2)),\n", " '--log-every', '25',\n", " '--device', DEVICE,\n", " ]\n", " print(' '.join(map(str, cmd)))\n", " subprocess.run(cmd, check=True)\n", "else:\n", " print('RUN_ROLEPLAY_SFT is False; skipping SFT.')\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 18. SFT report\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "if RUN_ROLEPLAY_SFT and (SFT_OUT_DIR / 'train_log.jsonl').exists():\n", " cmd = [\n", " sys.executable, '-m', 'ares_core.train_report',\n", " '--log', str(SFT_OUT_DIR / 'train_log.jsonl'),\n", " '--out', str(SFT_REPORT_PATH),\n", " '--title', f'Ares roleplay SFT report: {RUN_NAME}',\n", " ]\n", " subprocess.run(cmd, check=True)\n", " print('SFT report:', SFT_REPORT_PATH)\n", "else:\n", " print('No SFT report generated.')\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 19. Select final checkpoint\n", "\n", "The final checkpoint is the SFT best checkpoint if SFT ran, otherwise the base best checkpoint.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "if RUN_ROLEPLAY_SFT:\n", " FINAL_CKPT = SFT_OUT_DIR / 'ckpt_best.pt'\n", " if not FINAL_CKPT.exists():\n", " FINAL_CKPT = SFT_OUT_DIR / 'ckpt_last.pt'\n", "else:\n", " FINAL_CKPT = BASE_CKPT\n", "\n", "assert FINAL_CKPT.exists(), f'Final checkpoint missing: {FINAL_CKPT}'\n", "print('FINAL_CKPT =', FINAL_CKPT)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 20. REAL GENERATION: sample from Ares' trained brain\n", "\n", "This is the key cell. It loads the checkpoint and tokenizer and generates text from Ares weights.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompts = [\n", " 'Ares, explain why validation loss matters before scaling.',\n", " 'Roleplay as a machine-learning tutor in a lunar base and explain KV caches.',\n", " 'Ares, use Wikipedia-style knowledge to explain what a Transformer is.',\n", " 'Roleplay as a calm library archivist and help me plan a safe Ares training run.',\n", "]\n", "for prompt in prompts:\n", " print('\\n' + '='*100)\n", " print('PROMPT:', prompt)\n", " cmd = [\n", " sys.executable, '-m', 'ares_core.generate',\n", " '--checkpoint', str(FINAL_CKPT),\n", " '--tokenizer', str(TOKENIZER_PATH),\n", " '--prompt', prompt,\n", " '--max-new-tokens', '220',\n", " '--temperature', '0.75',\n", " '--top-k', '50',\n", " '--device', DEVICE,\n", " ]\n", " subprocess.run(cmd, check=True)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 21. Build SQLite RAG memory\n", "\n", "RAG is not the core brain, but it gives Ares a searchable memory store for grounding.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cmd = [sys.executable, '-m', 'ares_core.rag_sqlite', 'ingest', '--db', str(RAG_DB), '--input', str(TRAIN_PATH), str(ROLEPLAY_TRAIN)]\n", "print(' '.join(map(str, cmd)))\n", "subprocess.run(cmd, check=True)\n", "\n", "subprocess.run([\n", " sys.executable, '-m', 'ares_core.rag_sqlite', 'search',\n", " '--db', str(RAG_DB),\n", " '--query', 'validation loss roleplay transformer training',\n", " '--limit', '3',\n", "], check=True)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 22. REAL CHAT WRAPPER: Ares checkpoint + planning + RAG\n", "\n", "This cell calls `ares_core.agent_chat` using your trained checkpoint. It does not call an external AI model.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cmd = [\n", " sys.executable, '-m', 'ares_core.agent_chat',\n", " '--checkpoint', str(FINAL_CKPT),\n", " '--tokenizer', str(TOKENIZER_PATH),\n", " '--rag-db', str(RAG_DB),\n", " '--plan',\n", " '--prompt', 'Roleplay as a calm ML tutor and explain the Ares training loop clearly.',\n", " '--max-new-tokens', '260',\n", " '--temperature', '0.75',\n", " '--device', DEVICE,\n", "]\n", "subprocess.run(cmd, check=True)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 23. Optional interactive Ares chat\n", "\n", "Uncomment this after training if you want a simple Colab terminal chat.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Uncomment to chat interactively in the Colab output.\n", "# subprocess.run([\n", "# sys.executable, '-m', 'ares_core.agent_chat',\n", "# '--checkpoint', str(FINAL_CKPT),\n", "# '--tokenizer', str(TOKENIZER_PATH),\n", "# '--rag-db', str(RAG_DB),\n", "# '--plan',\n", "# '--interactive',\n", "# '--device', DEVICE,\n", "# ], check=True)\n", "print('Interactive chat command is ready. Uncomment it when desired.')\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 24. Save run summary\n", "\n", "This summary makes it easier to track experiments.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "summary = {\n", " 'run_name': RUN_NAME,\n", " 'profile': PROFILE,\n", " 'config_path': str(CONFIG_PATH),\n", " 'tokenizer_path': str(TOKENIZER_PATH),\n", " 'base_checkpoint_dir': str(OUT_DIR),\n", " 'sft_checkpoint_dir': str(SFT_OUT_DIR),\n", " 'final_checkpoint': str(FINAL_CKPT),\n", " 'rag_db': str(RAG_DB),\n", " 'wikipedia': f'{WIKIPEDIA_DATASET}:{WIKIPEDIA_CONFIG}',\n", " 'wikipedia_records': WIKIPEDIA_RECORDS,\n", " 'extra_hf_dataset': EXTRA_HF_DATASET,\n", " 'extra_hf_records': EXTRA_HF_RECORDS,\n", " 'ml_process_records': ML_PROCESS_RECORDS,\n", " 'roleplay_records': ROLEPLAY_RECORDS,\n", " 'train_steps': TRAIN_STEPS,\n", " 'roleplay_sft': RUN_ROLEPLAY_SFT,\n", " 'roleplay_sft_steps': ROLEPLAY_SFT_STEPS if RUN_ROLEPLAY_SFT else 0,\n", " 'created_at': time.strftime('%Y-%m-%d %H:%M:%S'),\n", "}\n", "SUMMARY_PATH.write_text(json.dumps(summary, indent=2), encoding='utf-8')\n", "print(SUMMARY_PATH)\n", "print(json.dumps(summary, indent=2))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 25. Optional upload to Hugging Face model repo\n", "\n", "Keep this disabled until you intentionally want to publish the trained checkpoint.\n", "\n", "Use `notebook_login()` or Colab secrets. Do **not** paste tokens into notebook cells that you will commit publicly.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "PUSH_TO_HUB = False\n", "HF_MODEL_REPO = 'jacmor64/ares-official-30m-colab'\n", "\n", "if PUSH_TO_HUB:\n", " from huggingface_hub import notebook_login, HfApi, upload_folder\n", " notebook_login()\n", " api = HfApi()\n", " api.create_repo(HF_MODEL_REPO, repo_type='model', exist_ok=True)\n", " upload_folder(folder_path=str(SFT_OUT_DIR if RUN_ROLEPLAY_SFT else OUT_DIR), repo_id=HF_MODEL_REPO, repo_type='model', path_in_repo='checkpoint')\n", " upload_folder(folder_path=str(TOKENIZER_PATH.parent), repo_id=HF_MODEL_REPO, repo_type='model', path_in_repo='tokenizer')\n", " upload_folder(folder_path=str(SUMMARY_PATH.parent), repo_id=HF_MODEL_REPO, repo_type='model', path_in_repo='summaries')\n", " print('Uploaded to', HF_MODEL_REPO)\n", "else:\n", " print('PUSH_TO_HUB is False. Artifacts remain in Drive.')\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Troubleshooting / Re-iteration Checklist\n", "\n", "Run through this checklist if anything fails:\n", "\n", "## If CUDA runs out of memory\n", "\n", "1. Change `PROFILE` to `smoke_10m`, or\n", "2. Keep `official_30m_512` and reduce `BATCH_SIZE` to `1`, or\n", "3. Use `context_30m_512` before trying `context_30m_1024`, or\n", "4. Lower `ROLEPLAY_RECORDS`, `WIKIPEDIA_RECORDS`, and `TRAIN_STEPS` for the first pass.\n", "\n", "## If loss is not improving\n", "\n", "1. Confirm tokenizer trained successfully.\n", "2. Confirm validation split exists.\n", "3. Confirm `val_loss` is logged.\n", "4. Reduce learning rate, for example `LEARNING_RATE = 1e-4`.\n", "5. Train longer only if validation loss is still decreasing.\n", "\n", "## If generation is repetitive\n", "\n", "1. Train longer on more data.\n", "2. Add higher-quality SFT examples.\n", "3. Use lower temperature, e.g. `0.6`.\n", "4. Remember: 30M is still small. Quality requires data + time + evaluation.\n", "\n", "## If you want stronger Ares\n", "\n", "Do this order:\n", "\n", "1. Finish `official_30m_512`.\n", "2. Try `context_30m_1024`.\n", "3. Increase Wikipedia records.\n", "4. Increase roleplay records.\n", "5. Train for more steps.\n", "6. Only then try the 100M config.\n" ] } ], "metadata": { "accelerator": "GPU", "colab": { "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python", "version": "3.x" } }, "nbformat": 4, "nbformat_minor": 5 }