{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# fa_dep_news_lg / fa_core_news_lg / fa_dep_news_trf / fa_core_news_trf — Colab training\n",
    "\n",
    "Trains the `lg` (bigger floret vectors) and `trf` (fine-tuned transformer) tiers of the\n",
    "Persian `spacy-fa-pipeline` project on a Colab GPU. `sm`/`md` are already built on CPU\n",
    "locally — this notebook only adds the two tiers that need real GPU memory.\n",
    "\n",
    "**`trf` uses `HooshvareLab/roberta-fa-zwnj-base` (Apache-2.0), not ParsBERT** — ParsBERT's\n",
    "model card carries no explicit licence, which is disqualifying for a package meant to be\n",
    "redistributed. See `TODO.md` in the repo.\n",
    "\n",
    "Floret vector *training* itself (the actual `lg`-tier 200k-row Wikipedia+OSCAR table) is\n",
    "not part of this notebook — that happens elsewhere (CPU-days, `spacy-vectors-builder`).\n",
    "This notebook only trains spaCy pipelines against whatever floret wheel you upload in\n",
    "step 6.\n",
    "\n",
    "## Before you run this\n",
    "\n",
    "1. **Runtime -> Change runtime type -> GPU** (a 16 GB T4/A10 is plenty for a base-size\n",
    "   transformer; no need for A100).\n",
    "2. Have ready, to upload when asked:\n",
    "   - A zip of the repo's **source only** (`git archive -o repo.zip HEAD` from the repo\n",
    "     root -- this naturally excludes everything `.gitignore` excludes: `assets/ corpus/\n",
    "     training/ metrics/ packages/ .venv/`). The self-hosted Gitea remote is LAN-only and\n",
    "     unreachable from Colab, so this notebook cannot `git clone` it directly.\n",
    "   - An `lg`-tier floret wheel (`fa_floret-0.1.0-py3-none-any-*.whl`) once it's built\n",
    "     elsewhere. If you don't have one yet, upload whatever `md`-tier wheel you have as a\n",
    "     stand-in -- the run will still be valid, just not the final `lg` numbers.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Confirm the GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Upload the repo source\n",
    "\n",
    "Upload the `repo.zip` produced by `git archive -o repo.zip HEAD` (run locally, in the repo\n",
    "root, before starting this notebook).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "import zipfile, pathlib\n",
    "\n",
    "REPO = pathlib.Path(\"/content/repo\")\n",
    "REPO.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "uploaded = files.upload()\n",
    "(zip_name,) = uploaded.keys()\n",
    "with zipfile.ZipFile(zip_name) as z:\n",
    "    z.extractall(REPO)\n",
    "\n",
    "%cd {REPO}\n",
    "!ls\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Colab ships a CUDA-enabled torch already; spacy[transformers] pulls in spacy-transformers +\n",
    "# a matching transformers/tokenizers. Installing spacy[cuda-autodetect] too is cheap insurance\n",
    "# for the GPU allocator path (unlike the local 940MX box, this doesn't need a manual cupy[ctk]\n",
    "# CUDA-toolkit install -- Colab's base image already has the CUDA libs on the system path).\n",
    "!pip install -q -U pip\n",
    "!pip install -q \"spacy[transformers,cuda-autodetect]\" spacy-transformers spacy-lookups-data\n",
    "\n",
    "import spacy, torch, spacy_transformers\n",
    "print(\"spacy\", spacy.__version__)\n",
    "print(\"spacy-transformers\", spacy_transformers.__version__)\n",
    "print(\"torch\", torch.__version__, \"cuda available:\", torch.cuda.is_available())\n",
    "\n",
    "from thinc.api import prefer_gpu\n",
    "print(\"thinc prefer_gpu:\", prefer_gpu())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Download the UD_Persian-PerDT assets\n",
    "\n",
    "Same public GitHub URLs and checksums as `project.yml` -- no private infrastructure needed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hashlib, urllib.request, pathlib\n",
    "\n",
    "ASSETS = [\n",
    "    (\"assets/ud/fa_perdt-ud-train.conllu\",\n",
    "     \"https://raw.githubusercontent.com/UniversalDependencies/UD_Persian-PerDT/master/fa_perdt-ud-train.conllu\",\n",
    "     \"f5a8ba901a776b4fd1941ecadcc6d506\"),\n",
    "    (\"assets/ud/fa_perdt-ud-dev.conllu\",\n",
    "     \"https://raw.githubusercontent.com/UniversalDependencies/UD_Persian-PerDT/master/fa_perdt-ud-dev.conllu\",\n",
    "     \"f103020da7c1e917aafb8a8321f4cb84\"),\n",
    "    (\"assets/ud/fa_perdt-ud-test.conllu\",\n",
    "     \"https://raw.githubusercontent.com/UniversalDependencies/UD_Persian-PerDT/master/fa_perdt-ud-test.conllu\",\n",
    "     \"b62a66994cef2c50f7e524a1471102d8\"),\n",
    "    (\"assets/ud-ner/train_with_NER_tag.txt\",\n",
    "     \"https://raw.githubusercontent.com/UniversalDependencies/UD_Persian-PerDT/master/not-to-release/Dadegan%20with%20NER%20tag/train_with_NER_tag.txt\",\n",
    "     \"ecb96cf99b38bc485cac21d22914e413\"),\n",
    "    (\"assets/ud-ner/dev_with_NER_tag.txt\",\n",
    "     \"https://raw.githubusercontent.com/UniversalDependencies/UD_Persian-PerDT/master/not-to-release/Dadegan%20with%20NER%20tag/dev_with_NER_tag.txt\",\n",
    "     \"2a56ef7eb2e3732e221317af457d1c09\"),\n",
    "    (\"assets/ud-ner/test_with_NER_tag.txt\",\n",
    "     \"https://raw.githubusercontent.com/UniversalDependencies/UD_Persian-PerDT/master/not-to-release/Dadegan%20with%20NER%20tag/test_with_NER_tag.txt\",\n",
    "     \"6d80dd783527562c2ea5189f218a12b5\"),\n",
    "]\n",
    "\n",
    "for dest, url, checksum in ASSETS:\n",
    "    dest = pathlib.Path(dest)\n",
    "    dest.parent.mkdir(parents=True, exist_ok=True)\n",
    "    urllib.request.urlretrieve(url, dest)\n",
    "    got = hashlib.md5(dest.read_bytes()).hexdigest()\n",
    "    status = \"OK\" if got == checksum else f\"MISMATCH (got {got})\"\n",
    "    print(f\"{dest}: {status}\")\n",
    "    assert got == checksum, f\"checksum mismatch on {dest}\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Build the UD + NER corpora (mirrors `project.yml`'s `convert-ud`/`transfer-ner`/`convert-ner`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy convert assets/ud/fa_perdt-ud-train.conllu corpus/merged --converter conllu --n-sents 10 --merge-subtokens\n",
    "!python -m spacy convert assets/ud/fa_perdt-ud-dev.conllu   corpus/merged --converter conllu --n-sents 10 --merge-subtokens\n",
    "!python -m spacy convert assets/ud/fa_perdt-ud-test.conllu  corpus/merged --converter conllu --n-sents 10 --merge-subtokens\n",
    "\n",
    "!python scripts/transfer_perdt_ner.py --conllu-dir assets/ud --ner-dir assets/ud-ner --out corpus/perdt-ner-iob\n",
    "\n",
    "!python -m spacy convert corpus/perdt-ner-iob/train.txt corpus/perdt-ner --converter ner --n-sents 10 --lang fa\n",
    "!python -m spacy convert corpus/perdt-ner-iob/dev.txt   corpus/perdt-ner --converter ner --n-sents 10 --lang fa\n",
    "!python -m spacy convert corpus/perdt-ner-iob/test.txt  corpus/perdt-ner --converter ner --n-sents 10 --lang fa\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. `lg`-tier floret vectors\n",
    "\n",
    "Upload a floret wheel (`fa_floret-0.1.0-py3-none-any-*.whl`), built elsewhere. Use the real\n",
    "200k-row Wikipedia+OSCAR table if you have one; otherwise upload whatever `md`-tier wheel\n",
    "you have as a stand-in -- the run will still be valid, just not the final `lg` numbers.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "\n",
    "uploaded = files.upload()\n",
    "(floret_wheel,) = uploaded.keys()\n",
    "\n",
    "!python scripts/unpack_vectors.py {floret_wheel} assets/vectors/fa_floret_lg\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Train the `lg` tier\n",
    "\n",
    "Byte-identical to the `md` recipe (`configs/fa_dep_news_md.cfg` / `configs/fa_ner_md.cfg`) --\n",
    "only `--paths.vectors` changes, isolating the effect of the bigger table exactly the way\n",
    "`md` isolated the effect of adding vectors over `sm`. No new config file needed.\n",
    "\n",
    "`--gpu-id 0` for both here: the earlier CPU-vs-GPU timing experiment ran on a 2 GB GTX 940MX,\n",
    "where the small NER architecture's transfer/launch overhead beat its GPU compute win. A 16 GB\n",
    "Colab GPU has far more bandwidth/compute headroom, so that conclusion may not hold here --\n",
    "worth timing both `--gpu-id 0` and `--gpu-id -1` yourself if you want to confirm.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy train configs/fa_dep_news_md.cfg --output training/dep-lg \\\n",
    "    --paths.train corpus/merged/fa_perdt-ud-train.spacy \\\n",
    "    --paths.dev corpus/merged/fa_perdt-ud-dev.spacy \\\n",
    "    --paths.vectors assets/vectors/fa_floret_lg \\\n",
    "    --gpu-id 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy train configs/fa_ner_md.cfg --output training/perdt-ner-lg \\\n",
    "    --paths.train corpus/perdt-ner/train.spacy \\\n",
    "    --paths.dev corpus/perdt-ner/dev.spacy \\\n",
    "    --paths.vectors assets/vectors/fa_floret_lg \\\n",
    "    --gpu-id 0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Assemble + evaluate `lg`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy benchmark accuracy training/dep-lg/model-best corpus/merged/fa_perdt-ud-test.spacy \\\n",
    "    --output metrics/lg-ud-test.json --gpu-id 0\n",
    "\n",
    "!python scripts/finalize_pipeline.py training/dep-lg/model-best training/fa_dep_news_lg \\\n",
    "    --variant dep --size lg --version 3.8.0 --ud-metrics metrics/lg-ud-test.json\n",
    "\n",
    "!python scripts/finalize_pipeline.py training/dep-lg/model-best training/fa_core_news_lg \\\n",
    "    --variant core --size lg --version 3.8.0 --add-ner training/perdt-ner-lg/model-best\n",
    "\n",
    "!python -m spacy benchmark accuracy training/fa_core_news_lg corpus/merged/fa_perdt-ud-test.spacy \\\n",
    "    --output metrics/lg-core-ud-test.json --gpu-id 0\n",
    "!python -m spacy benchmark accuracy training/fa_core_news_lg corpus/perdt-ner/test.spacy \\\n",
    "    --output metrics/lg-perdt-ner-test.json --gpu-id 0\n",
    "\n",
    "!python scripts/finalize_pipeline.py training/dep-lg/model-best training/fa_core_news_lg \\\n",
    "    --variant core --size lg --version 3.8.0 --add-ner training/perdt-ner-lg/model-best \\\n",
    "    --ud-metrics metrics/lg-core-ud-test.json --ner-metrics metrics/lg-perdt-ner-test.json\n",
    "\n",
    "!python scripts/smoke_test.py training/fa_core_news_lg\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Generate the `trf` configs\n",
    "\n",
    "`spacy init config --optimize accuracy -G` fills in a valid `spacy-transformers`\n",
    "architecture automatically (letting spaCy own the schema instead of hand-writing one).\n",
    "The only edit afterward is swapping the default transformer name for\n",
    "`HooshvareLab/roberta-fa-zwnj-base` and turning on mixed precision, since 16 GB has room\n",
    "for it.\n",
    "\n",
    "Same split as `sm`/`md`/`lg`: `dep` (tagger/morphologizer/lemmatizer/parser) and `ner`\n",
    "trained as separate pipelines, each with its own transformer, so `ner` can be re-sourced\n",
    "into `core` afterward exactly like the CPU tiers.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy init config configs/fa_dep_news_trf.cfg --lang fa \\\n",
    "    --pipeline tagger,morphologizer,trainable_lemmatizer,parser \\\n",
    "    --optimize accuracy -G --force\n",
    "\n",
    "!python -m spacy init config configs/fa_ner_trf.cfg --lang fa \\\n",
    "    --pipeline ner \\\n",
    "    --optimize accuracy -G --force\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import re\n",
    "\n",
    "TRANSFORMER_NAME = \"HooshvareLab/roberta-fa-zwnj-base\"\n",
    "\n",
    "for path in [\"configs/fa_dep_news_trf.cfg\", \"configs/fa_ner_trf.cfg\"]:\n",
    "    text = open(path, encoding=\"utf8\").read()\n",
    "    # Swap whatever default transformer `init config` picked for roberta-fa-zwnj-base.\n",
    "    text = re.sub(\n",
    "        r'(\\[components\\.transformer\\.model\\]\\nname = )\"[^\"]+\"',\n",
    "        lambda m: m.group(1) + '\"' + TRANSFORMER_NAME + '\"',\n",
    "        text,\n",
    "    )\n",
    "    # 16 GB has room for mixed precision; halves activation memory, meaningfully faster.\n",
    "    if \"mixed_precision\" in text:\n",
    "        text = text.replace(\"mixed_precision = false\", \"mixed_precision = true\")\n",
    "    else:\n",
    "        text = text.replace(\"[training]\\n\", \"[training]\\nmixed_precision = true\\n\", 1)\n",
    "    open(path, \"w\", encoding=\"utf8\").write(text)\n",
    "    print(\"patched\", path, \"transformer =\", TRANSFORMER_NAME)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy debug config configs/fa_dep_news_trf.cfg \\\n",
    "    --paths.train corpus/merged/fa_perdt-ud-train.spacy \\\n",
    "    --paths.dev corpus/merged/fa_perdt-ud-dev.spacy\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Train the `trf` tier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy train configs/fa_dep_news_trf.cfg --output training/dep-trf \\\n",
    "    --paths.train corpus/merged/fa_perdt-ud-train.spacy \\\n",
    "    --paths.dev corpus/merged/fa_perdt-ud-dev.spacy \\\n",
    "    --gpu-id 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy train configs/fa_ner_trf.cfg --output training/perdt-ner-trf \\\n",
    "    --paths.train corpus/perdt-ner/train.spacy \\\n",
    "    --paths.dev corpus/perdt-ner/dev.spacy \\\n",
    "    --gpu-id 0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Assemble + evaluate `trf`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy benchmark accuracy training/dep-trf/model-best corpus/merged/fa_perdt-ud-test.spacy \\\n",
    "    --output metrics/trf-ud-test.json --gpu-id 0\n",
    "\n",
    "!python scripts/finalize_pipeline.py training/dep-trf/model-best training/fa_dep_news_trf \\\n",
    "    --variant dep --size trf --version 3.8.0 --ud-metrics metrics/trf-ud-test.json\n",
    "\n",
    "!python scripts/finalize_pipeline.py training/dep-trf/model-best training/fa_core_news_trf \\\n",
    "    --variant core --size trf --version 3.8.0 --add-ner training/perdt-ner-trf/model-best\n",
    "\n",
    "!python -m spacy benchmark accuracy training/fa_core_news_trf corpus/merged/fa_perdt-ud-test.spacy \\\n",
    "    --output metrics/trf-core-ud-test.json --gpu-id 0\n",
    "!python -m spacy benchmark accuracy training/fa_core_news_trf corpus/perdt-ner/test.spacy \\\n",
    "    --output metrics/trf-perdt-ner-test.json --gpu-id 0\n",
    "\n",
    "!python scripts/finalize_pipeline.py training/dep-trf/model-best training/fa_core_news_trf \\\n",
    "    --variant core --size trf --version 3.8.0 --add-ner training/perdt-ner-trf/model-best \\\n",
    "    --ud-metrics metrics/trf-core-ud-test.json --ner-metrics metrics/trf-perdt-ner-test.json\n",
    "\n",
    "!python scripts/smoke_test.py training/fa_core_news_trf\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 12. Compare every tier\n",
    "\n",
    "Reads whichever `metrics/*-ud-test.json` / `metrics/*-perdt-ner-test.json` files exist in\n",
    "this Colab session (only `lg` and `trf`, produced above). To compare against the local\n",
    "`sm`/`md` numbers, upload `metrics/core-ud-test.json`, `metrics/perdt-ner-test.json`,\n",
    "`metrics/md-core-ud-test.json`, `metrics/md-perdt-ner-test.json` from the repo first.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, pathlib\n",
    "\n",
    "ROWS = [\n",
    "    (\"sm\",  \"metrics/core-ud-test.json\",     \"metrics/perdt-ner-test.json\"),\n",
    "    (\"md\",  \"metrics/md-core-ud-test.json\",  \"metrics/md-perdt-ner-test.json\"),\n",
    "    (\"lg\",  \"metrics/lg-core-ud-test.json\",  \"metrics/lg-perdt-ner-test.json\"),\n",
    "    (\"trf\", \"metrics/trf-core-ud-test.json\", \"metrics/trf-perdt-ner-test.json\"),\n",
    "]\n",
    "\n",
    "def load(path):\n",
    "    p = pathlib.Path(path)\n",
    "    return json.loads(p.read_text()) if p.exists() else None\n",
    "\n",
    "print(f\"{'tier':<5}{'tag_acc':>9}{'dep_las':>9}{'lemma_acc':>11}{'ents_f':>9}\")\n",
    "for tier, ud_path, ner_path in ROWS:\n",
    "    ud, ner = load(ud_path), load(ner_path)\n",
    "    tag = f\"{ud['tag_acc']*100:.2f}\" if ud and ud.get('tag_acc') is not None else \"-\"\n",
    "    las = f\"{ud['dep_las']*100:.2f}\" if ud and ud.get('dep_las') is not None else \"-\"\n",
    "    lem = f\"{ud['lemma_acc']*100:.2f}\" if ud and ud.get('lemma_acc') is not None else \"-\"\n",
    "    entf = f\"{ner['ents_f']*100:.2f}\" if ner and ner.get('ents_f') is not None else \"-\"\n",
    "    print(f\"{tier:<5}{tag:>9}{las:>9}{lem:>11}{entf:>9}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 13. Download the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "\n",
    "!zip -r /content/lg_trf_results.zip training/fa_dep_news_lg training/fa_core_news_lg \\\n",
    "    training/fa_dep_news_trf training/fa_core_news_trf metrics assets/vectors/fa_floret_lg \\\n",
    "    configs/fa_dep_news_trf.cfg configs/fa_ner_trf.cfg\n",
    "\n",
    "files.download(\"/content/lg_trf_results.zip\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 14. (Optional) package as installable wheels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy package training/fa_dep_news_lg packages --name dep_news_lg --version 3.8.0 --build sdist,wheel --force\n",
    "!python -m spacy package training/fa_core_news_lg packages --name core_news_lg --version 3.8.0 --build sdist,wheel --force\n",
    "!python -m spacy package training/fa_dep_news_trf packages --name dep_news_trf --version 3.8.0 --build sdist,wheel --force\n",
    "!python -m spacy package training/fa_core_news_trf packages --name core_news_trf --version 3.8.0 --build sdist,wheel --force\n",
    "\n",
    "!zip -r /content/packages.zip packages\n",
    "files.download(\"/content/packages.zip\")\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "fa_lg_trf_training.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}