Enhance Colab badge functionality: Update script to add 'Open in Colab' badges to both template and solution notebooks. Modify badge generation to accommodate different notebook directories and ensure proper processing of both templates and solutions.

Author: duoan

scripts/add_colab_badges.py

@@ -1,27 +1,29 @@
 #!/usr/bin/env python3
-"""Add 'Open in Colab' badges to all template notebooks."""
+"""Add 'Open in Colab' badges to all template and solution notebooks."""
 
 import json
 from pathlib import Path
 
 REPO = "duoan/TorchCode"
 BRANCH = "master"
-TEMPLATES_DIR = Path(__file__).resolve().parent.parent / "templates"
+ROOT = Path(__file__).resolve().parent.parent
+TEMPLATES_DIR = ROOT / "templates"
+SOLUTIONS_DIR = ROOT / "solutions"
 BADGE_IMG = "https://colab.research.google.com/assets/colab-badge.svg"
 
 
-def colab_url(filename: str) -> str:
+def colab_url(filename: str, folder: str) -> str:
     return (
         f"https://colab.research.google.com/github/{REPO}"
-        f"/blob/{BRANCH}/templates/{filename}"
+        f"/blob/{BRANCH}/{folder}/{filename}"
     )
 
 
-def badge_markdown(filename: str) -> str:
-    return f"[![Open In Colab]({BADGE_IMG})]({colab_url(filename)})"
+def badge_markdown(filename: str, folder: str) -> str:
+    return f"[![Open In Colab]({BADGE_IMG})]({colab_url(filename, folder)})"
 
 
-def process_notebook(path: Path) -> bool:
+def process_notebook(path: Path, folder: str) -> bool:
     with open(path, "r", encoding="utf-8") as f:
         nb = json.load(f)
 
@@ -34,7 +36,7 @@ def process_notebook(path: Path) -> bool:
     if "colab-badge.svg" in flat:
         return False
 
-    badge = badge_markdown(path.name)
+    badge = badge_markdown(path.name, folder)
     cells[0]["source"] = [badge + "\n\n"] + (
         source_lines if isinstance(source_lines, list) else [source_lines]
     )
@@ -49,11 +51,17 @@ def process_notebook(path: Path) -> bool:
 def main() -> None:
     updated = 0
     for nb_path in sorted(TEMPLATES_DIR.glob("*.ipynb")):
-        if process_notebook(nb_path):
-            print(f"  ✅ {nb_path.name}")
+        if process_notebook(nb_path, "templates"):
+            print(f"  ✅ templates/{nb_path.name}")
             updated += 1
         else:
-            print(f"  ⏭️  {nb_path.name} (already has badge or skipped)")
+            print(f"  ⏭️  templates/{nb_path.name} (already has badge or skipped)")
+    for nb_path in sorted(SOLUTIONS_DIR.glob("*.ipynb")):
+        if process_notebook(nb_path, "solutions"):
+            print(f"  ✅ solutions/{nb_path.name}")
+            updated += 1
+        else:
+            print(f"  ⏭️  solutions/{nb_path.name} (already has badge or skipped)")
     print(f"\nDone — updated {updated} notebooks.")
 
 

scripts/add_colab_torch_judge_install.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Add Colab-only pip install of torch-judge to all notebooks that use torch_judge."""
+
+import json
+from pathlib import Path
+
+INSTALL_CELL_SOURCE = [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n",
+]
+
+MARKER = "get_ipython().run_line_magic('pip', 'install"
+
+
+def has_torch_judge(nb: dict) -> bool:
+    for cell in nb.get("cells", []):
+        src = cell.get("source", [])
+        flat = "".join(src) if isinstance(src, list) else str(src)
+        if "torch_judge" in flat:
+            return True
+    return False
+
+
+def already_has_install(nb: dict) -> bool:
+    for cell in nb.get("cells", []):
+        src = cell.get("source", [])
+        flat = "".join(src) if isinstance(src, list) else str(src)
+        if MARKER in flat and "torch-judge" in flat:
+            return True
+    return False
+
+
+def process_notebook(path: Path) -> bool:
+    with open(path, "r", encoding="utf-8") as f:
+        nb = json.load(f)
+
+    if not has_torch_judge(nb):
+        return False
+    if already_has_install(nb):
+        return False
+
+    cells = nb["cells"]
+    if not cells:
+        return False
+
+    # Insert install cell at index 1 (after first cell, usually markdown title)
+    install_cell = {
+        "cell_type": "code",
+        "metadata": {},
+        "source": INSTALL_CELL_SOURCE,
+        "outputs": [],
+        "execution_count": None,
+    }
+    cells.insert(1, install_cell)
+
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(nb, f, ensure_ascii=False, indent=1)
+        f.write("\n")
+
+    return True
+
+
+def main() -> None:
+    root = Path(__file__).resolve().parent.parent
+    updated = 0
+    for pattern in ["templates/*.ipynb", "solutions/*.ipynb"]:
+        for path in sorted(root.glob(pattern)):
+            if process_notebook(path):
+                print(f"  + {path.relative_to(root)}")
+                updated += 1
+    print(f"Updated {updated} notebooks.")
+
+
+if __name__ == "__main__":
+    main()

solutions/01_relu_solution.ipynb

@@ -1,73 +1,89 @@
 {
- "nbformat": 4,
- "nbformat_minor": 5,
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.11.0"
-  }
- },
  "cells": [
   {
    "cell_type": "markdown",
+   "id": "0556419b",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/01_relu_solution.ipynb)\n",
+    "\n",
     "# 🟢 Solution: Implement ReLU\n",
     "\n",
     "Reference solution for the ReLU activation function.\n",
     "\n",
     "$$\\text{ReLU}(x) = \\max(0, x)$$"
-   ],
-   "outputs": []
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
-    "import torch"
-   ],
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
    "outputs": [],
-   "execution_count": null
+   "source": [
+    "import torch"
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "# ✅ SOLUTION\n",
     "\n",
     "def relu(x: torch.Tensor) -> torch.Tensor:\n",
     "    return x * (x > 0).float()"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "# Verify\n",
     "x = torch.tensor([-2., -1., 0., 1., 2.])\n",
     "print(\"Input: \", x)\n",
     "print(\"Output:\", relu(x))"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "# Run judge\n",
     "from torch_judge import check\n",
     "check(\"relu\")"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   }
- ]
-}
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

solutions/02_softmax_solution.ipynb

@@ -17,6 +17,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/02_softmax_solution.ipynb)\n\n",
     "# 🟢 Solution: Implement Softmax\n",
     "\n",
     "Reference solution for the numerically-stable Softmax function.\n",
@@ -25,6 +26,20 @@
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -73,4 +88,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/03_linear_solution.ipynb

@@ -17,12 +17,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/03_linear_solution.ipynb)\n\n",
     "# 🟡 Solution: Simple Linear Layer\n",
     "\n",
     "Reference solution for a fully-connected linear layer: **y = xW^T + b**"
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -77,4 +92,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/04_layernorm_solution.ipynb

@@ -17,6 +17,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/04_layernorm_solution.ipynb)\n\n",
     "# 🟡 Solution: Implement LayerNorm\n",
     "\n",
     "Reference solution for Layer Normalization.\n",
@@ -25,6 +26,20 @@
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -76,4 +91,4 @@
    "execution_count": null
   }
  ]
-}
+}