{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Togyzkumalak OCR — Colab training (moves + diagram)\n",
    "\n",
    "Runtime → Change runtime type → **GPU**, then run the cells in order.\n",
    "\n",
    "Trains both classifiers on on-the-fly synthetic data (nothing to upload):\n",
    "1. **moves** — 163-class move-cell reader (`11`..`99x` + `empty`)\n",
    "2. **diagram** — 85-class board-diagram reader (values `0`-`81`, `x`, `-`, `empty`) for the kazan boxes and 2×9 pit grids of the summary strips\n",
    "\n",
    "At the end you download one `artifacts.zip` with the ONNX models + class sidecars, ready to drop into `models/` of the HuggingFace Space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch the code at runtime.\n",
    "# Option A (default): clone from GitHub\n",
    "!git clone https://github.com/ansarzeinulla/9OCR.git project\n",
    "%cd project\n",
    "\n",
    "# Option B: upload a zip named project.zip instead, then:\n",
    "# from google.colab import files; files.upload()\n",
    "# !unzip -q project.zip -d project && cd project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q -r requirements-train.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build glyph pools at runtime: ARDIS (European handwriting) + EMNIST via\n",
    "# torchvision (--emnist is needed on Colab because ingredients/ is gitignored;\n",
    "# digit 0 comes from these pools).\n",
    "!python scripts/get_glyphs.py --emnist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sanity check both synthesizers before burning GPU time.\n",
    "!python -m togyz.synth --preview preview_moves.png --task moves\n",
    "!python -m togyz.synth --preview preview_diagram.png --task diagram\n",
    "from IPython.display import Image as IPImage, display\n",
    "display(IPImage('preview_moves.png'))\n",
    "display(IPImage('preview_diagram.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the move classifier (~30 epochs x 50k synthetic cells).\n",
    "!python train.py --task moves --epochs 30 --samples-per-epoch 50000 --batch-size 256 --workers 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the unified board-diagram classifier.\n",
    "!python train.py --task diagram --epochs 30 --samples-per-epoch 50000 --batch-size 256 --workers 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export both checkpoints to single-file ONNX (+ .classes.json sidecars).\n",
    "!python scripts/export_onnx.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bundle the serving artifacts and download.\n",
    "# In the Space's models/ dir: best.onnx stays best.onnx, the diagram model\n",
    "# is renamed to diagram.onnx (+ diagram.classes.json).\n",
    "import shutil, zipfile\n",
    "from pathlib import Path\n",
    "\n",
    "out = Path('artifacts'); out.mkdir(exist_ok=True)\n",
    "shutil.copy('checkpoints/best.onnx', out / 'best.onnx')\n",
    "shutil.copy('checkpoints/best.classes.json', out / 'best.classes.json')\n",
    "shutil.copy('checkpoints/diagram/best.onnx', out / 'diagram.onnx')\n",
    "shutil.copy('checkpoints/diagram/best.classes.json', out / 'diagram.classes.json')\n",
    "with zipfile.ZipFile('artifacts.zip', 'w') as zf:\n",
    "    for f in out.iterdir():\n",
    "        zf.write(f, arcname=f.name)\n",
    "\n",
    "from google.colab import files\n",
    "files.download('artifacts.zip')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {"provenance": []},
  "kernelspec": {"display_name": "Python 3", "name": "python3"},
  "language_info": {"name": "python"}
 },
 "nbformat": 4,
 "nbformat_minor": 5
}