{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Ares Complete Brain Notebook: Wikipedia + ML Processes + Roleplay\n",
        "\n",
        "This notebook is the ready-to-run Colab path for making Ares generate responses from **its own trained checkpoint**.\n",
        "\n",
        "It builds:\n",
        "\n",
        "1. A streamed Wikipedia subset from Hugging Face.\n",
        "2. A deterministic machine-learning-process curriculum.\n",
        "3. A standalone roleplay JSONL dataset with thousands of scenarios.\n",
        "4. A mixed pretraining corpus.\n",
        "5. A held-out validation set.\n",
        "6. A BPE tokenizer trained from scratch.\n",
        "7. A small Ares Transformer checkpoint trained from scratch.\n",
        "8. Optional roleplay-focused SFT from the trained checkpoint.\n",
        "9. Generation and chat cells that sample from Ares weights, not an external AI API.\n",
        "\n",
        "**Important:** The Hugging Face Static Space cannot run GPU training or server-side inference. This notebook runs in Colab. Keep checkpoints in Google Drive or a separate model repo, not the Static Space repo.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 1. Runtime check\n",
        "\n",
        "Set **Runtime \u2192 Change runtime type \u2192 GPU** before running.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os, sys, subprocess, json, time\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: GPU is not enabled. Use Runtime \u2192 Change runtime type \u2192 GPU.')\n",
        "except Exception as e:\n",
        "    print('Torch check failed:', repr(e))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 2. Install dependencies\n",
        "\n",
        "Colab usually already includes a CUDA PyTorch build, so this avoids reinstalling torch.\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 Ares project\n",
        "\n",
        "Use either:\n",
        "\n",
        "- Upload `ares-static-space.zip` to Colab, or\n",
        "- Clone the Hugging Face Static Space repo.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "import os, sys, subprocess, zipfile\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 project zip...')\n",
        "        with zipfile.ZipFile(ZIP_PATH) as z:\n",
        "            z.extractall('/content')\n",
        "    else:\n",
        "        print('Cloning from Hugging Face Space...')\n",
        "        subprocess.run(['git', 'clone', '--depth', '1', REPO_URL, str(PROJECT)], check=True)\n",
        "\n",
        "assert (PROJECT / 'ares_core').exists(), 'Ares project not found.'\n",
        "os.chdir(PROJECT)\n",
        "sys.path.insert(0, str(PROJECT))\n",
        "print('PROJECT =', PROJECT)\n",
        "print('Ares files ready.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 4. Store artifacts in Google Drive\n",
        "\n",
        "This prevents Colab runtime resets from deleting checkpoints. Do **not** commit these artifacts to the public Static Space repo.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "USE_DRIVE = True\n",
        "if USE_DRIVE:\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/drive')\n",
        "    ARTIFACTS = Path('/content/drive/MyDrive/ares_brain_artifacts')\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. Configure the complete brain-training run\n",
        "\n",
        "For the first run, keep `USE_FULL_EN_WIKIPEDIA = False`. That uses Simple English Wikipedia for a manageable smoke/full pipeline. Switch to full English Wikipedia later when you have more time/storage.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "RUN_NAME = 'ares-complete-wiki-roleplay-30m'\n",
        "\n",
        "# Model config. If you hit CUDA OOM, change to configs/ares_8m.json.\n",
        "CONFIG_PATH = PROJECT / 'configs' / 'ares_30m.json'\n",
        "\n",
        "# Wikipedia source.\n",
        "USE_FULL_EN_WIKIPEDIA = False\n",
        "WIKIPEDIA_DATASET = 'wikimedia/wikipedia'\n",
        "WIKIPEDIA_CONFIG = '20231101.en' if USE_FULL_EN_WIKIPEDIA else '20231101.simple'\n",
        "WIKIPEDIA_RECORDS = 8000        # raise later: 25k, 100k, ...\n",
        "\n",
        "# Extra training curriculum.\n",
        "ML_PROCESS_RECORDS = 4000\n",
        "ROLEPLAY_RECORDS = 12000        # thousands of roleplay scenarios\n",
        "ROLEPLAY_SEED = 2026\n",
        "\n",
        "# Tokenizer and training.\n",
        "TOKENIZER_VOCAB_SIZE = 8192\n",
        "TRAIN_STEPS = 1000              # smoke: 300; better: 2k-10k+\n",
        "BATCH_SIZE = 4                  # lower on OOM\n",
        "GRAD_ACCUM = 2\n",
        "LEARNING_RATE = 3e-4\n",
        "EVAL_EVERY = 100\n",
        "EVAL_BATCHES = 30\n",
        "SAVE_EVERY = 200\n",
        "DEVICE = 'auto'\n",
        "\n",
        "# Optional roleplay-focused SFT after mixed training.\n",
        "RUN_ROLEPLAY_SFT = True\n",
        "ROLEPLAY_SFT_STEPS = 300\n",
        "ROLEPLAY_SFT_LR = 1e-4\n",
        "\n",
        "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",
        "for p in [DATA_DIR, TOKENIZER_PATH.parent, OUT_DIR, SFT_OUT_DIR, RAG_DB.parent]:\n",
        "    p.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "print(json.dumps({\n",
        "    'run': RUN_NAME,\n",
        "    'config': str(CONFIG_PATH),\n",
        "    'wikipedia': f'{WIKIPEDIA_DATASET}:{WIKIPEDIA_CONFIG}',\n",
        "    'wiki_records': WIKIPEDIA_RECORDS,\n",
        "    'ml_records': ML_PROCESS_RECORDS,\n",
        "    'roleplay_records': ROLEPLAY_RECORDS,\n",
        "    'train_steps': TRAIN_STEPS,\n",
        "    'roleplay_sft': RUN_ROLEPLAY_SFT,\n",
        "    'artifacts': str(ARTIFACTS),\n",
        "}, indent=2))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 6. Generate standalone roleplay JSONL dataset\n",
        "\n",
        "This creates the explicit roleplay dataset you asked for. It is deterministic and generated locally, not by 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', str(ROLEPLAY_SEED),\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('Preview:')\n",
        "for i, line in enumerate(ROLEPLAY_JSONL.open('r', encoding='utf-8')):\n",
        "    print(line[:600])\n",
        "    if i >= 2:\n",
        "        break\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 7. Convert roleplay JSONL to SFT text\n",
        "\n",
        "This gives Ares a chat-formatted roleplay dataset for optional supervised fine-tuning.\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')[:1200])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 8. Build Wikipedia + ML + roleplay mixed corpus\n",
        "\n",
        "This streams Wikipedia from Hugging Face. It does not download the full dump unless you set huge record counts.\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', str(ROLEPLAY_SEED),\n",
        "    '--min-chars', '60',\n",
        "    '--max-chars', '24000',\n",
        "]\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.read_text()[:2000])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 9. Split mixed corpus into train/validation\n",
        "\n",
        "Validation tells us whether Ares is actually improving instead of memorizing.\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": [
        "## 10. Split roleplay SFT data\n",
        "\n",
        "This is used for the optional roleplay refinement pass.\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": [
        "## 11. Train Ares tokenizer from scratch\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "cmd = [\n",
        "    sys.executable, '-m', 'ares_core.tokenizer_train',\n",
        "    '--input', str(TRAIN_PATH),\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": [
        "## 12. Check parameter counts\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": [
        "## 13. Train Ares' mixed brain\n",
        "\n",
        "This is the main from-scratch training run on Wikipedia + ML processes + roleplay.\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",
        "    '--eval-every', str(EVAL_EVERY),\n",
        "    '--eval-batches', str(EVAL_BATCHES),\n",
        "    '--save-every', str(SAVE_EVERY),\n",
        "    '--device', DEVICE,\n",
        "]\n",
        "print(' '.join(map(str, cmd)))\n",
        "subprocess.run(cmd, check=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 14. Generate mixed-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 brain 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('Display skipped:', e)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 15. Optional roleplay-focused SFT\n",
        "\n",
        "This resumes from the best mixed checkpoint, resets optimizer/step count, and trains on roleplay chat data.\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 checkpoint:', 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",
        "        '--eval-every', str(max(50, ROLEPLAY_SFT_STEPS // 3)),\n",
        "        '--eval-batches', str(EVAL_BATCHES),\n",
        "        '--save-every', str(max(100, ROLEPLAY_SFT_STEPS // 2)),\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.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 16. Generate SFT report if enabled\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": [
        "## 17. Choose final 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",
        "print('FINAL_CKPT =', FINAL_CKPT)\n",
        "assert FINAL_CKPT.exists(), 'Final checkpoint not found.'\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 18. Generate using Ares' trained brain\n",
        "\n",
        "This loads your checkpoint and samples tokens from Ares. No external AI model is called.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "prompts = [\n",
        "    'Ares, explain how validation loss prevents overfitting.',\n",
        "    'Roleplay as a machine-learning tutor in a lunar base and explain KV caches.',\n",
        "    'Ares, summarize what Wikipedia-style knowledge helps you learn.',\n",
        "    'Roleplay as a calm library archivist and help me plan a safe Transformer training run.',\n",
        "]\n",
        "for prompt in prompts:\n",
        "    print('\n",
        "' + '='*80)\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', '160',\n",
        "        '--temperature', '0.8',\n",
        "        '--top-k', '50',\n",
        "        '--device', DEVICE,\n",
        "    ]\n",
        "    subprocess.run(cmd, check=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 19. Build RAG memory from the training corpus\n",
        "\n",
        "RAG is not the core brain, but it gives Ares retrievable memory.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "subprocess.run([sys.executable, '-m', 'ares_core.rag_sqlite', 'ingest', '--db', str(RAG_DB), '--input', str(TRAIN_PATH)], check=True)\n",
        "subprocess.run([sys.executable, '-m', 'ares_core.rag_sqlite', 'search', '--db', str(RAG_DB), '--query', 'machine learning validation loss roleplay', '--limit', '3'], check=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 20. Chat wrapper: brain + planning + RAG\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 why Ares starts small before scaling.',\n",
        "    '--max-new-tokens', '220',\n",
        "    '--temperature', '0.75',\n",
        "    '--device', DEVICE,\n",
        "]\n",
        "subprocess.run(cmd, check=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 21. Optional interactive chat\n",
        "\n",
        "Run this if you want to talk with Ares in the Colab cell output.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Uncomment to start interactive mode.\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 cell is ready; uncomment it when desired.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 22. Optional: publish the trained brain to a model repo\n",
        "\n",
        "Do **not** put checkpoints in the Static Space repo. Use a separate Hugging Face model repo if you decide to publish weights.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "PUSH_TO_HUB = False\n",
        "HF_MODEL_REPO = 'jacmor64/ares-complete-wiki-roleplay-30m'\n",
        "\n",
        "if PUSH_TO_HUB:\n",
        "    from huggingface_hub import notebook_login, HfApi, upload_folder\n",
        "    notebook_login()  # use Colab login/secrets; do not paste tokens into repo files\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",
        "    print('Uploaded to', HF_MODEL_REPO)\n",
        "else:\n",
        "    print('PUSH_TO_HUB is False. Checkpoints remain in Drive.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 23. Scaling checklist\n",
        "\n",
        "If the output is weak, do not jump straight to 1B. Improve in this order:\n",
        "\n",
        "1. Confirm train and validation loss decrease.\n",
        "2. Increase `WIKIPEDIA_RECORDS`, `ML_PROCESS_RECORDS`, and `ROLEPLAY_RECORDS`.\n",
        "3. Increase `TRAIN_STEPS`.\n",
        "4. Use better curated roleplay/SFT examples.\n",
        "5. Move from 10M/30M to 100M only when data and loss curves are stable.\n",
        "6. Use `20231101.en` for larger English Wikipedia runs.\n",
        "7. Use a multi-GPU setup for 1B-class configs.\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
}