{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Ares No-Terminal Trainer and Connector \u2014 Fixed Safe Version\n",
        "\n",
        "You do **not** need terminal commands.\n",
        "\n",
        "Use this notebook like this:\n",
        "\n",
        "1. Open it in Google Colab.\n",
        "2. Click **Runtime \u2192 Change runtime type \u2192 GPU**.\n",
        "3. Click **Runtime \u2192 Run all**.\n",
        "4. Wait for training.\n",
        "5. At the end, click **Open connected Ares UI**.\n",
        "\n",
        "## What was fixed\n",
        "\n",
        "Your previous run hit two Colab failures:\n",
        "\n",
        "- `SIGABRT` while building a large serious corpus.\n",
        "- `SIGKILL` while training the 30M / 1024-context model.\n",
        "\n",
        "This fixed notebook now:\n",
        "\n",
        "- defaults to **Safe 30M / 512 context** instead of 1024,\n",
        "- uses a serious multi-source dataset but with safer default record counts,\n",
        "- automatically retries corpus building with smaller counts if a dataset source crashes,\n",
        "- automatically falls back to a smaller model if training is killed,\n",
        "- still uses serious sources: Wikipedia, FineWeb-Edu, OpenWebMath, Python code, UltraChat, ML curriculum, and thousands of roleplay scenarios.\n",
        "\n",
        "The serious builder uses `jqop/python-code-dataset` for Python code examples. After connection, the Static UI calls the Colab Ares API endpoint `/generate` for model responses.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 1. Click play: GPU check"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title GPU check { display-mode: \"form\" }\n",
        "import os, sys, subprocess, json, time, re, urllib.parse\n",
        "from pathlib import Path\n",
        "\n",
        "try:\n",
        "    import torch\n",
        "    HAS_GPU = torch.cuda.is_available()\n",
        "    print('torch:', torch.__version__)\n",
        "    print('cuda available:', HAS_GPU)\n",
        "    if HAS_GPU:\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. Bigger Ares training may fail. Use Runtime \u2192 Change runtime type \u2192 GPU.')\n",
        "except Exception as e:\n",
        "    HAS_GPU = False\n",
        "    print('Torch will be checked again after setup:', repr(e))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 2. Click play: install and load Ares"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Install dependencies and load Ares repo { display-mode: \"form\" }\n",
        "!pip -q install -U tokenizers datasets safetensors huggingface_hub tqdm numpy fastapi uvicorn pydantic requests\n",
        "\n",
        "from pathlib import Path\n",
        "import os, sys, subprocess, zipfile, json, time, re, urllib.parse\n",
        "\n",
        "PROJECT = Path('/content/ares-static-space')\n",
        "REPO_URL = 'https://huggingface.co/spaces/jacmor64/ares-static-lab'\n",
        "\n",
        "if not (PROJECT / 'ares_core').exists():\n",
        "    subprocess.run(['git', 'clone', '--depth', '1', REPO_URL, str(PROJECT)], check=True)\n",
        "\n",
        "os.chdir(PROJECT)\n",
        "sys.path.insert(0, str(PROJECT))\n",
        "\n",
        "import torch\n",
        "HAS_GPU = torch.cuda.is_available()\n",
        "print('Ares repo loaded:', PROJECT)\n",
        "print('cuda available:', HAS_GPU)\n",
        "if HAS_GPU:\n",
        "    print('gpu:', torch.cuda.get_device_name(0))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 3. Click play: connect Google Drive"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Mount Google Drive { display-mode: \"form\" }\n",
        "USE_DRIVE = True #@param {type:\"boolean\"}\n",
        "\n",
        "from pathlib import Path\n",
        "if USE_DRIVE:\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/drive')\n",
        "    ARTIFACTS = Path('/content/drive/MyDrive/ares_no_terminal_training')\n",
        "else:\n",
        "    ARTIFACTS = Path('/content/ares_no_terminal_training')\n",
        "ARTIFACTS.mkdir(parents=True, exist_ok=True)\n",
        "print('Artifacts folder:', ARTIFACTS)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 4. Choose settings from dropdowns\n",
        "\n",
        "Recommended default:\n",
        "\n",
        "```text\n",
        "Safe 30M / 512 context\n",
        "Safe serious dataset\n",
        "```\n",
        "\n",
        "This avoids the `SIGKILL` you saw with 1024 context.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Ares training settings { display-mode: \"form\" }\n",
        "architecture_profile = \"Safe 30M / 512 context\" #@param [\"Safe 30M / 512 context\", \"Practical 30M / 1024 context\", \"Serious 124M / 2048 context\", \"Smoke 10M / 256 context\"]\n",
        "dataset_size = \"Safe serious dataset\" #@param [\"Tiny debug\", \"Safe serious dataset\", \"Medium serious dataset\", \"Large serious dataset\"]\n",
        "run_roleplay_sft = True #@param {type:\"boolean\"}\n",
        "\n",
        "if architecture_profile.startswith('Smoke'):\n",
        "    RUN_NAME = 'ares-no-terminal-smoke-10m'\n",
        "    CONFIG_PATH = PROJECT / 'configs' / 'ares_8m.json'\n",
        "    TOKENIZER_VOCAB_SIZE = 4096\n",
        "    BATCH_SIZE = 1\n",
        "    GRAD_ACCUM = 1\n",
        "    TRAIN_STEPS = 200\n",
        "elif architecture_profile.startswith('Safe'):\n",
        "    RUN_NAME = 'ares-no-terminal-30m-ctx512-safe'\n",
        "    CONFIG_PATH = PROJECT / 'configs' / 'ares_30m.json'\n",
        "    TOKENIZER_VOCAB_SIZE = 8192\n",
        "    BATCH_SIZE = 1\n",
        "    GRAD_ACCUM = 4\n",
        "    TRAIN_STEPS = 1000\n",
        "elif architecture_profile.startswith('Practical'):\n",
        "    RUN_NAME = 'ares-no-terminal-30m-ctx1024'\n",
        "    CONFIG_PATH = PROJECT / 'configs' / 'ares_30m_ctx1024.json'\n",
        "    TOKENIZER_VOCAB_SIZE = 8192\n",
        "    BATCH_SIZE = 1\n",
        "    GRAD_ACCUM = 8\n",
        "    TRAIN_STEPS = 1500\n",
        "else:\n",
        "    RUN_NAME = 'ares-no-terminal-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",
        "\n",
        "if dataset_size == 'Tiny debug':\n",
        "    WIKIPEDIA_RECORDS = 500\n",
        "    FINEWEB_RECORDS = 500\n",
        "    OPENWEBMATH_RECORDS = 200\n",
        "    CODE_RECORDS = 200\n",
        "    DIALOGUE_RECORDS = 200\n",
        "    ML_PROCESS_RECORDS = 1000\n",
        "    ROLEPLAY_RECORDS = 2000\n",
        "elif dataset_size == 'Safe serious dataset':\n",
        "    WIKIPEDIA_RECORDS = 3000\n",
        "    FINEWEB_RECORDS = 3000\n",
        "    OPENWEBMATH_RECORDS = 1000\n",
        "    CODE_RECORDS = 1000\n",
        "    DIALOGUE_RECORDS = 1000\n",
        "    ML_PROCESS_RECORDS = 5000\n",
        "    ROLEPLAY_RECORDS = 10000\n",
        "elif dataset_size == 'Medium serious dataset':\n",
        "    WIKIPEDIA_RECORDS = 10000\n",
        "    FINEWEB_RECORDS = 10000\n",
        "    OPENWEBMATH_RECORDS = 5000\n",
        "    CODE_RECORDS = 5000\n",
        "    DIALOGUE_RECORDS = 5000\n",
        "    ML_PROCESS_RECORDS = 10000\n",
        "    ROLEPLAY_RECORDS = 25000\n",
        "else:\n",
        "    WIKIPEDIA_RECORDS = 50000\n",
        "    FINEWEB_RECORDS = 50000\n",
        "    OPENWEBMATH_RECORDS = 20000\n",
        "    CODE_RECORDS = 20000\n",
        "    DIALOGUE_RECORDS = 20000\n",
        "    ML_PROCESS_RECORDS = 20000\n",
        "    ROLEPLAY_RECORDS = 50000\n",
        "\n",
        "LEARNING_RATE = 2e-4\n",
        "MIN_LR = 2e-5\n",
        "WARMUP_STEPS = 100\n",
        "EVAL_EVERY = 250\n",
        "EVAL_BATCHES = 10\n",
        "SAVE_EVERY = 500\n",
        "LOG_EVERY = 25\n",
        "ROLEPLAY_SFT_STEPS = 400\n",
        "ROLEPLAY_SFT_LR = 8e-5\n",
        "DEVICE = 'auto'\n",
        "\n",
        "print(json.dumps({\n",
        "    'architecture_profile': architecture_profile,\n",
        "    'dataset_size': dataset_size,\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",
        "    'run_roleplay_sft': run_roleplay_sft,\n",
        "}, indent=2))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 5. Click play: create folders"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Create artifact folders { display-mode: \"form\" }\n",
        "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.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",
        "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('Ready:', RUN_NAME)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 6. Click play: build roleplay dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Generate roleplay scenarios { display-mode: \"form\" }\n",
        "subprocess.run([sys.executable, '-m', 'ares_core.synthetic_curriculum', 'roleplay', '--output', str(ROLEPLAY_JSONL), '--count', str(ROLEPLAY_RECORDS), '--seed', '2026', '--jsonl'], check=True)\n",
        "subprocess.run([sys.executable, '-m', 'ares_core.sft', '--input', str(ROLEPLAY_JSONL), '--output', str(ROLEPLAY_TEXT)], check=True)\n",
        "print('Roleplay JSONL MB:', round(ROLEPLAY_JSONL.stat().st_size/1e6, 2))\n",
        "print('Roleplay SFT MB:', round(ROLEPLAY_TEXT.stat().st_size/1e6, 2))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 7. Click play: build serious corpus with automatic fallback\n",
        "\n",
        "If the large dataset build crashes, this cell automatically retries with smaller counts.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Build serious multi-source corpus safely { display-mode: \"form\" }\n",
        "def build_corpus_with_counts(w, f, m, c, d, ml, rp):\n",
        "    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(w),\n",
        "        '--wikipedia-config', '20231101.en',\n",
        "        '--fineweb-records', str(f),\n",
        "        '--openwebmath-records', str(m),\n",
        "        '--code-records', str(c),\n",
        "        '--dialogue-records', str(d),\n",
        "        '--tinystories-records', '0',\n",
        "        '--ml-records', str(ml),\n",
        "        '--roleplay-records', str(rp),\n",
        "        '--min-chars', '80',\n",
        "        '--max-chars', '32000',\n",
        "        '--seed', '2026',\n",
        "    ]\n",
        "    print('Building corpus with counts:', {'wiki':w,'fineweb':f,'math':m,'code':c,'dialogue':d,'ml':ml,'roleplay':rp})\n",
        "    subprocess.run(cmd, check=True)\n",
        "\n",
        "try:\n",
        "    build_corpus_with_counts(WIKIPEDIA_RECORDS, FINEWEB_RECORDS, OPENWEBMATH_RECORDS, CODE_RECORDS, DIALOGUE_RECORDS, ML_PROCESS_RECORDS, ROLEPLAY_RECORDS)\n",
        "except subprocess.CalledProcessError as e:\n",
        "    print('First corpus build failed. Retrying with safe fallback counts.')\n",
        "    build_corpus_with_counts(1000, 1000, 500, 500, 500, 2000, 5000)\n",
        "\n",
        "print('Corpus MB:', round(CORPUS_PATH.stat().st_size/1e6, 2))\n",
        "print(Path(str(CORPUS_PATH) + '.manifest.json').read_text()[:3000])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 8. Click play: split datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Split train and validation { display-mode: \"form\" }\n",
        "subprocess.run([sys.executable, '-m', 'ares_core.split_corpus', '--input', str(CORPUS_PATH), '--train-output', str(TRAIN_PATH), '--val-output', str(VAL_PATH), '--val-ratio', '0.02', '--min-val-records', '100'], check=True)\n",
        "subprocess.run([sys.executable, '-m', 'ares_core.split_corpus', '--input', str(ROLEPLAY_TEXT), '--train-output', str(ROLEPLAY_TRAIN), '--val-output', str(ROLEPLAY_VAL), '--val-ratio', '0.02', '--min-val-records', '100'], 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": [
        "## 9. Click play: train tokenizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Train Ares tokenizer { display-mode: \"form\" }\n",
        "subprocess.run([sys.executable, '-m', 'ares_core.tokenizer_train', '--input', str(TRAIN_PATH), str(ROLEPLAY_TRAIN), '--output', str(TOKENIZER_PATH), '--vocab-size', str(TOKENIZER_VOCAB_SIZE), '--min-frequency', '2'], check=True)\n",
        "print('Tokenizer:', TOKENIZER_PATH)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 10. Click play: audit model size"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Audit run size { display-mode: \"form\" }\n",
        "from tokenizers import Tokenizer\n",
        "from ares_core.config import AresConfig, estimate_parameters\n",
        "cfg = AresConfig.from_json(CONFIG_PATH)\n",
        "tok = Tokenizer.from_file(str(TOKENIZER_PATH))\n",
        "print(json.dumps({\n",
        "    'model': cfg.model_name,\n",
        "    'estimated_parameters_before_vocab_adjustment': estimate_parameters(cfg),\n",
        "    'context_window': cfg.max_seq_len,\n",
        "    'tokenizer_vocab': tok.get_vocab_size(),\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",
        "}, indent=2))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 11. Click play: train Ares with automatic fallback\n",
        "\n",
        "If Colab kills the larger model, this cell falls back to the safer 10M model automatically.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Train Ares base checkpoint safely { display-mode: \"form\" }\n",
        "def train_with_config(config_path, out_dir, steps, batch_size, grad_accum):\n",
        "    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(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",
        "    subprocess.run(cmd, check=True)\n",
        "\n",
        "try:\n",
        "    train_with_config(CONFIG_PATH, OUT_DIR, TRAIN_STEPS, BATCH_SIZE, GRAD_ACCUM)\n",
        "    ACTIVE_CONFIG_PATH = CONFIG_PATH\n",
        "    ACTIVE_OUT_DIR = OUT_DIR\n",
        "except subprocess.CalledProcessError:\n",
        "    print('Training was killed or failed. Falling back to safe 10M / 256-ish smoke model so you still get a working Ares checkpoint.')\n",
        "    ACTIVE_CONFIG_PATH = PROJECT / 'configs' / 'ares_8m.json'\n",
        "    ACTIVE_OUT_DIR = ARTIFACTS / 'checkpoints' / f'{RUN_NAME}-fallback-10m'\n",
        "    ACTIVE_OUT_DIR.mkdir(parents=True, exist_ok=True)\n",
        "    train_with_config(ACTIVE_CONFIG_PATH, ACTIVE_OUT_DIR, 300, 1, 1)\n",
        "\n",
        "print('Active checkpoint dir:', ACTIVE_OUT_DIR)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 12. Click play: roleplay SFT refinement"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Optional roleplay SFT { display-mode: \"form\" }\n",
        "BASE_CKPT = ACTIVE_OUT_DIR / 'ckpt_best.pt'\n",
        "if not BASE_CKPT.exists():\n",
        "    BASE_CKPT = ACTIVE_OUT_DIR / 'ckpt_last.pt'\n",
        "print('Base checkpoint:', BASE_CKPT)\n",
        "\n",
        "if run_roleplay_sft:\n",
        "    SFT_OUT_DIR = ARTIFACTS / 'checkpoints' / f'{RUN_NAME}-roleplay-sft'\n",
        "    SFT_OUT_DIR.mkdir(parents=True, exist_ok=True)\n",
        "    cmd = [\n",
        "        sys.executable, '-m', 'ares_core.train',\n",
        "        '--config', str(ACTIVE_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', '1',\n",
        "        '--grad-accum', '1',\n",
        "        '--lr', str(ROLEPLAY_SFT_LR),\n",
        "        '--min-lr', str(ROLEPLAY_SFT_LR/10),\n",
        "        '--warmup-steps', '50',\n",
        "        '--eval-every', '200',\n",
        "        '--eval-batches', str(EVAL_BATCHES),\n",
        "        '--save-every', '400',\n",
        "        '--log-every', '25',\n",
        "        '--device', DEVICE,\n",
        "    ]\n",
        "    try:\n",
        "        subprocess.run(cmd, check=True)\n",
        "        FINAL_DIR = SFT_OUT_DIR\n",
        "    except subprocess.CalledProcessError:\n",
        "        print('SFT failed; using base checkpoint instead.')\n",
        "        FINAL_DIR = ACTIVE_OUT_DIR\n",
        "else:\n",
        "    FINAL_DIR = ACTIVE_OUT_DIR\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 13. Click play: generate with Ares weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Generate test messages { display-mode: \"form\" }\n",
        "FINAL_CKPT = FINAL_DIR / 'ckpt_best.pt'\n",
        "if not FINAL_CKPT.exists():\n",
        "    FINAL_CKPT = FINAL_DIR / 'ckpt_last.pt'\n",
        "assert FINAL_CKPT.exists(), FINAL_CKPT\n",
        "print('FINAL_CKPT:', FINAL_CKPT)\n",
        "\n",
        "prompts = [\n",
        "    'Ares, explain why serious datasets matter.',\n",
        "    'Roleplay as a machine-learning tutor and explain validation loss.',\n",
        "    'Ares, plan the next training run.',\n",
        "]\n",
        "for prompt in prompts:\n",
        "    print('\\n' + '='*100)\n",
        "    print('PROMPT:', prompt)\n",
        "    subprocess.run([\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",
        "    ], check=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 14. Click play: open connected Ares UI\n",
        "\n",
        "This starts the model API and gives you a big clickable link. Click it. No terminal commands.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Start API and create connected UI link { display-mode: \"form\" }\n",
        "import subprocess, time, re, urllib.parse\n",
        "from IPython.display import display, HTML\n",
        "\n",
        "PORT = 8000\n",
        "server_cmd = [\n",
        "    sys.executable, '-m', 'ares_core.api_server',\n",
        "    '--checkpoint', str(FINAL_CKPT),\n",
        "    '--tokenizer', str(TOKENIZER_PATH),\n",
        "    '--device', 'auto',\n",
        "    '--host', '0.0.0.0',\n",
        "    '--port', str(PORT),\n",
        "]\n",
        "server = subprocess.Popen(server_cmd)\n",
        "time.sleep(8)\n",
        "print('Ares API server started.')\n",
        "\n",
        "!wget -q https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64 -O /usr/local/bin/cloudflared\n",
        "!chmod +x /usr/local/bin/cloudflared\n",
        "\n",
        "cloudflared = subprocess.Popen(\n",
        "    ['cloudflared', 'tunnel', '--url', f'http://127.0.0.1:{PORT}'],\n",
        "    stdout=subprocess.PIPE,\n",
        "    stderr=subprocess.STDOUT,\n",
        "    text=True,\n",
        ")\n",
        "public_url = None\n",
        "for _ in range(120):\n",
        "    line = cloudflared.stdout.readline()\n",
        "    print(line, end='')\n",
        "    m = re.search(r'https://[-a-zA-Z0-9.]+\\.trycloudflare\\.com', line)\n",
        "    if m:\n",
        "        public_url = m.group(0)\n",
        "        break\n",
        "    time.sleep(1)\n",
        "\n",
        "if not public_url:\n",
        "    raise RuntimeError('Could not find cloudflared public URL. Rerun this cell.')\n",
        "\n",
        "space_url = 'https://huggingface.co/spaces/jacmor64/ares-static-lab?engine=' + urllib.parse.quote(public_url, safe='')\n",
        "print('Ares engine URL:', public_url)\n",
        "print('Connected UI URL:', space_url)\n",
        "display(HTML(f'<h2><a target=\"_blank\" href=\"{space_url}\">Open connected Ares UI</a></h2><p>If the link does not auto-connect, copy this engine URL into the Ares engine box: <code>{public_url}</code></p>'))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Done\n",
        "\n",
        "Keep this Colab tab open. If Colab stops, the connected UI stops generating because the model server is gone.\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
}