Add COLAB routine

examples/prompt-tuning-personachat.ipynb

@@ -29,17 +29,26 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8ec7e625",
+   "id": "73bbc648",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# uncomment this block for colab usage\n",
-    "# git clone https://github.com/bigscience-workshop/petals\n",
-    "# !pip install petals/requirements.txt\n",
-    "# !pip install datasets wandb\n",
+    "# This block is only need for colab users. It will change nothing if you are running this notebook locally.\n",
+    "import subprocess\n",
+    "import sys\n",
+    "\n",
+    "\n",
+    "IN_COLAB = 'google.colab' in sys.modules\n",
+    "\n",
+    "if IN_COLAB:\n",
+    "    subprocess.run(['git', 'clone', 'https://github.com/bigscience-workshop/petals'])\n",
+    "    subprocess.run(['pip', 'install', '-r', 'petals/requirements.txt'])\n",
+    "    subprocess.run(['pip', 'install', 'datasets', 'lib64'])\n",
     "\n",
-    "# if you use cpu-only setup on colab, please uncomment this block\n",
-    "# !rm -r /usr/local/cuda/lib64"
+    "    try:\n",
+    "        subprocess.check_output([\"nvidia-smi\", \"-L\"])\n",
+    "    except subprocess.CalledProcessError as e:\n",
+    "        subprocess.run(['rm', '-r', '/usr/local/cuda/lib64'])"
    ]
   },
   {
@@ -317,7 +326,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.8.9"
   },
   "vscode": {
    "interpreter": {