{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# OFFICIAL Ares Serious Dataset Training Run\n",
        "\n",
        "This notebook replaces the toy/starter loop with a serious, multi-source training loop for Ares.\n",
        "\n",
        "The previous 1,000-roleplay / 500-ML-process dataset is **not** the training corpus. It is only a downloadable preview and smoke-test artifact. This notebook builds the real corpus in Colab.\n",
        "\n",
        "## Serious default data mixture\n",
        "\n",
        "The default serious profile uses:\n",
        "\n",
        "- `wikimedia/wikipedia` \u2014 full English Wikipedia by default.\n",
        "- `HuggingFaceFW/fineweb-edu` \u2014 high-quality educational web text sample.\n",
        "- `open-web-math/open-web-math` \u2014 mathematical text.\n",
        "- `jqop/python-code-dataset` \u2014 code/documentation data.\n",
        "- `HuggingFaceH4/ultrachat_200k` \u2014 multi-turn dialogue/instruction data.\n",
        "- Ares ML-process curriculum \u2014 generated locally.\n",
        "- Ares roleplay curriculum \u2014 **50,000 roleplay scenarios by default**.\n",
        "\n",
        "## Serious default architecture\n",
        "\n",
        "The default architecture is:\n",
        "\n",
        "```text\n",
        "configs/ares_124m.json\n",
        "~100M parameters\n",
        "2048-token context\n",
        "32K tokenizer vocabulary\n",
        "```\n",
        "\n",
        "This is not a toy seed model. It is still small compared to frontier LLMs, but it is a serious Colab-scale Ares training target.\n",
        "\n",
        "## Hardware expectation\n",
        "\n",
        "A free T4 Colab may struggle with the 124M/2048 profile. If it OOMs, use `PROFILE = \"practical_30m_1024\"` while keeping the serious datasets. Do not go back to the starter dataset as the main corpus.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 1. Enable GPU\n",
        "\n",
        "In Colab:\n",
        "\n",
        "```text\n",
        "Runtime \u2192 Change runtime type \u2192 GPU\n",
        "```\n",
        "\n",
        "Run this cell and confirm CUDA is true.\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: CUDA is false. Enable GPU before real training.')\n",
        "except Exception as e:\n",
        "    print('Torch check failed before install:', repr(e))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 2. Install dependencies\n",
        "\n",
        "This does not intentionally replace Colab's CUDA PyTorch 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. Clone the official Ares 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",
        "        with zipfile.ZipFile(ZIP_PATH) as z:\n",
        "            z.extractall('/content')\n",
        "    else:\n",
        "        subprocess.run(['git', 'clone', '--depth', '1', REPO_URL, str(PROJECT)], check=True)\n",
        "\n",
        "assert (PROJECT / 'ares_core').exists(), 'Ares repo not found.'\n",
        "os.chdir(PROJECT)\n",
        "sys.path.insert(0, str(PROJECT))\n",
        "print('PROJECT =', PROJECT)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 4. Mount Google Drive\n",
        "\n",
        "Checkpoints and datasets go to Drive.\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_serious_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 serious architecture profile\n",
        "\n",
        "Default:\n",
        "\n",
        "```python\n",
        "PROFILE = \"serious_124m_2048\"\n",
        "```\n",
        "\n",
        "If Colab runs out of memory, switch to:\n",
        "\n",
        "```python\n",
        "PROFILE = \"practical_30m_1024\"\n",
        "```\n",
        "\n",
        "The data remains serious either way.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "PROFILE = 'serious_124m_2048'  # 'serious_124m_2048' or 'practical_30m_1024'\n",
        "\n",
        "if PROFILE == 'serious_124m_2048':\n",
        "    RUN_NAME = 'ares-serious-124m-ctx2048'\n",
        "    CONFIG_PATH = PROJECT / 'configs' / 'ares_124m.json'\n",
        "    TOKENIZER_VOCAB_SIZE = 32000\n",
        "    BATCH_SIZE = 1\n",
        "    GRAD_ACCUM = 8\n",
        "    TRAIN_STEPS = 3000\n",
        "elif PROFILE == 'practical_30m_1024':\n",
        "    RUN_NAME = 'ares-serious-30m-ctx1024'\n",
        "    CONFIG_PATH = PROJECT / 'configs' / 'ares_30m_ctx1024.json'\n",
        "    TOKENIZER_VOCAB_SIZE = 32000\n",
        "    BATCH_SIZE = 1\n",
        "    GRAD_ACCUM = 8\n",
        "    TRAIN_STEPS = 3000\n",
        "else:\n",
        "    raise ValueError(PROFILE)\n",
        "\n",
        "# Serious corpus scale. Raise these later; lower only if you must debug.\n",
        "WIKIPEDIA_RECORDS = 50_000\n",
        "WIKIPEDIA_CONFIG = '20231101.en'\n",
        "FINEWEB_RECORDS = 50_000\n",
        "OPENWEBMATH_RECORDS = 20_000\n",
        "CODE_RECORDS = 20_000\n",
        "DIALOGUE_RECORDS = 20_000\n",
        "TINYSTORIES_RECORDS = 0  # optional; leave 0 for serious run\n",
        "ML_PROCESS_RECORDS = 20_000\n",
        "ROLEPLAY_RECORDS = 50_000\n",
        "\n",
        "LEARNING_RATE = 2e-4\n",
        "MIN_LR = 2e-5\n",
        "WARMUP_STEPS = 200\n",
        "EVAL_EVERY = 250\n",
        "EVAL_BATCHES = 20\n",
        "SAVE_EVERY = 500\n",
        "LOG_EVERY = 25\n",
        "DEVICE = 'auto'\n",
        "RUN_ROLEPLAY_SFT = True\n",
        "ROLEPLAY_SFT_STEPS = 800\n",
        "ROLEPLAY_SFT_LR = 8e-5\n",
        "\n",
        "print(json.dumps({\n",
        "    'profile': PROFILE,\n",
        "    'run_name': RUN_NAME,\n",
        "    'config': str(CONFIG_PATH),\n",
        "    'tokenizer_vocab': TOKENIZER_VOCAB_SIZE,\n",
        "    'wikipedia_records': WIKIPEDIA_RECORDS,\n",
        "    'fineweb_records': FINEWEB_RECORDS,\n",
        "    'openwebmath_records': OPENWEBMATH_RECORDS,\n",
        "    'code_records': CODE_RECORDS,\n",
        "    'dialogue_records': DIALOGUE_RECORDS,\n",
        "    'ml_process_records': ML_PROCESS_RECORDS,\n",
        "    'roleplay_records': ROLEPLAY_RECORDS,\n",
        "    'train_steps': TRAIN_STEPS,\n",
        "}, indent=2))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 6. Artifact paths"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "DATA_DIR = ARTIFACTS / 'data' / RUN_NAME\n",
        "CORPUS_PATH = DATA_DIR / 'serious_corpus.txt'\n",
        "TRAIN_PATH = DATA_DIR / 'serious_train.txt'\n",
        "VAL_PATH = DATA_DIR / 'serious_val.txt'\n",
        "ROLEPLAY_JSONL = DATA_DIR / 'roleplay_dataset_50000.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",
        "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",
        "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",
        "print('DATA_DIR:', DATA_DIR)\n",
        "print('OUT_DIR:', OUT_DIR)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 7. Generate 50,000 Ares roleplay scenarios\n",
        "\n",
        "This is not the old 1,000-example starter. This is a large local roleplay curriculum.\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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 8. Convert roleplay JSONL to SFT text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "cmd = [sys.executable, '-m', 'ares_core.sft', '--input', str(ROLEPLAY_JSONL), '--output', str(ROLEPLAY_TEXT)]\n",
        "print(' '.join(map(str, cmd)))\n",
        "subprocess.run(cmd, check=True)\n",
        "print('Roleplay SFT MB:', round(ROLEPLAY_TEXT.stat().st_size / 1e6, 2))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 9. Build serious multi-dataset corpus\n",
        "\n",
        "This streams capped records from official/public datasets and combines them with Ares curriculum.\n",
        "\n",
        "If a source fails due to a temporary Hugging Face issue, the builder records the error in the manifest and continues. Check the manifest after this cell.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "cmd = [\n",
        "    sys.executable, '-m', 'ares_core.serious_mixture_build',\n",
        "    '--output', str(CORPUS_PATH),\n",
        "    '--local', str(PROJECT / 'data' / 'sample_corpus.txt'),\n",
        "    '--preset', 'serious',\n",
        "    '--wikipedia-records', str(WIKIPEDIA_RECORDS),\n",
        "    '--wikipedia-config', WIKIPEDIA_CONFIG,\n",
        "    '--fineweb-records', str(FINEWEB_RECORDS),\n",
        "    '--openwebmath-records', str(OPENWEBMATH_RECORDS),\n",
        "    '--code-records', str(CODE_RECORDS),\n",
        "    '--dialogue-records', str(DIALOGUE_RECORDS),\n",
        "    '--tinystories-records', str(TINYSTORIES_RECORDS),\n",
        "    '--ml-records', str(ML_PROCESS_RECORDS),\n",
        "    '--roleplay-records', str(ROLEPLAY_RECORDS),\n",
        "    '--min-chars', '80',\n",
        "    '--max-chars', '32000',\n",
        "    '--seed', '2026',\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 = Path(str(CORPUS_PATH) + '.manifest.json')\n",
        "print(manifest_path.read_text()[:4000])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 10. Split train/validation\n",
        "\n",
        "Never skip 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', '1000',\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"
      ]
    },
    {
      "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', '1000',\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 32K BPE tokenizer from scratch\n",
        "\n",
        "This tokenizer is trained on the serious corpus, not on the starter preview.\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. Token and parameter audit"
      ]
    },
    {
      "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 token_count(path):\n",
        "    text = Path(path).read_text(encoding='utf-8', errors='ignore')\n",
        "    return len(tok.encode(text).ids)\n",
        "\n",
        "stats = {\n",
        "    'config': str(CONFIG_PATH),\n",
        "    'model_name': cfg.model_name,\n",
        "    'estimated_parameters': estimate_parameters(cfg),\n",
        "    'context_window': cfg.max_seq_len,\n",
        "    'tokenizer_vocab': tok.get_vocab_size(),\n",
        "    'train_tokens': token_count(TRAIN_PATH),\n",
        "    'val_tokens': token_count(VAL_PATH),\n",
        "    'roleplay_sft_tokens': token_count(ROLEPLAY_TRAIN),\n",
        "    'effective_tokens_per_step': cfg.max_seq_len * BATCH_SIZE * GRAD_ACCUM,\n",
        "    'planned_tokens_seen': cfg.max_seq_len * BATCH_SIZE * GRAD_ACCUM * TRAIN_STEPS,\n",
        "}\n",
        "print(json.dumps(stats, indent=2))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 14. Train Ares base model\n",
        "\n",
        "This is the serious base pretraining loop. Watch `val_loss`, not just training loss.\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": [
        "## 15. Base training report"
      ]
    },
    {
      "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 serious base 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 display skipped:', e)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 16. Roleplay SFT refinement\n",
        "\n",
        "This improves chat/roleplay behavior after base training.\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', '100',\n",
        "        '--eval-every', str(max(100, ROLEPLAY_SFT_STEPS // 4)),\n",
        "        '--eval-batches', str(EVAL_BATCHES),\n",
        "        '--save-every', str(max(200, 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('SFT skipped.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 17. SFT report"
      ]
    },
    {
      "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 serious 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.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 18. Select final checkpoint"
      ]
    },
    {
      "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",
        "assert FINAL_CKPT.exists(), FINAL_CKPT\n",
        "print('FINAL_CKPT:', FINAL_CKPT)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 19. REAL GENERATION from Ares weights\n",
        "\n",
        "This is the proof cell. It loads your checkpoint and samples from Ares. No external AI API is used.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "prompts = [\n",
        "    'Ares, explain why a serious dataset is required for general intelligence.',\n",
        "    'Roleplay as a machine-learning tutor and explain validation loss clearly.',\n",
        "    'Use Wikipedia-style knowledge to explain what a Transformer model does.',\n",
        "    'Ares, plan the next training run after this one.',\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', '260',\n",
        "        '--temperature', '0.75',\n",
        "        '--top-k', '50',\n",
        "        '--device', DEVICE,\n",
        "    ]\n",
        "    subprocess.run(cmd, check=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 20. Build RAG memory from serious corpus"
      ]
    },
    {
      "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",
        "subprocess.run([\n",
        "    sys.executable, '-m', 'ares_core.rag_sqlite', 'search',\n",
        "    '--db', str(RAG_DB),\n",
        "    '--query', 'Transformer validation loss general intelligence roleplay',\n",
        "    '--limit', '5',\n",
        "], check=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 21. Ares chat wrapper: checkpoint + planning + RAG"
      ]
    },
    {
      "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', 'Ares, explain the serious training loop and why starter data is not enough.',\n",
        "    '--max-new-tokens', '320',\n",
        "    '--temperature', '0.75',\n",
        "    '--device', DEVICE,\n",
        "]\n",
        "subprocess.run(cmd, check=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 22. Save summary"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "summary = {\n",
        "    'run_name': RUN_NAME,\n",
        "    'profile': PROFILE,\n",
        "    'config': str(CONFIG_PATH),\n",
        "    'final_checkpoint': str(FINAL_CKPT),\n",
        "    'tokenizer': str(TOKENIZER_PATH),\n",
        "    'rag_db': str(RAG_DB),\n",
        "    'wikipedia_records': WIKIPEDIA_RECORDS,\n",
        "    'fineweb_records': FINEWEB_RECORDS,\n",
        "    'openwebmath_records': OPENWEBMATH_RECORDS,\n",
        "    'code_records': CODE_RECORDS,\n",
        "    'dialogue_records': DIALOGUE_RECORDS,\n",
        "    'ml_process_records': ML_PROCESS_RECORDS,\n",
        "    'roleplay_records': ROLEPLAY_RECORDS,\n",
        "    'train_steps': TRAIN_STEPS,\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": [
        "## 23. Optional upload to Hugging Face model repo\n",
        "\n",
        "Leave this off until you intentionally want to publish the trained weights.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "PUSH_TO_HUB = False\n",
        "HF_MODEL_REPO = 'jacmor64/ares-serious-124m'\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",
        "else:\n",
        "    print('PUSH_TO_HUB is False. Weights remain in Drive.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Re-iteration checklist\n",
        "\n",
        "If this fails:\n",
        "\n",
        "1. Do **not** reduce to the starter dataset as the main corpus.\n",
        "2. If GPU OOM: switch from `serious_124m_2048` to `practical_30m_1024`.\n",
        "3. If a dataset source fails: lower that source count or temporarily set it to `0`, but keep multiple serious sources.\n",
        "4. If validation loss rises: stop, lower learning rate, add data, or shorten SFT.\n",
        "5. If generation is repetitive: train longer on more records and improve SFT quality.\n",
        "6. The next serious scaling step is more data first, then 124M longer, then larger 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
}