Update Querying Data to run directly from Colab

Querying Data.ipynb

@@ -33,10 +33,11 @@
     "import pandas.io.formats.style\n",
     "import random\n",
     "import functools\n",
+    "import os\n",
     "from typing import Literal\n",
     "\n",
     "SOURCE: Literal[\"danbooru\", \"e621\"] = \"e621\"\n",
-    "DATA_FOLDER = \"H:/Data/TagSuggest/e621_dataframes\""
+    "DATA_FOLDER = \".\""
    ]
   },
   {
@@ -105,10 +106,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tags = pandas.read_feather(f\"{DATA_FOLDER}/tags.feather\")\n",
-    "posts_by_tag = pandas.read_feather(f\"{DATA_FOLDER}/posts_by_tag.feather\").set_index(\"tag_id\")\n",
-    "tags_by_post = pandas.read_feather(f\"{DATA_FOLDER}/tags_by_post.feather\").set_index(\"post_id\")\n",
-    "implications = pandas.read_feather(f\"{DATA_FOLDER}/implications.feather\")\n",
+    "def get_feather(filename: str) -> pandas.DataFrame:\n",
+    "    target = f\"{DATA_FOLDER}/{filename}.feather\"\n",
+    "    if not os.path.exists(target):\n",
+    "        !wget -O {filename}.feather https://huggingface.co/Specimen5423/E621TagAssociations/resolve/main/{filename}.feather?download=true\n",
+    "    return pandas.read_feather(target)\n",
+    "\n",
+    "os.makedirs(DATA_FOLDER, exist_ok=True)\n",
+    "tags = get_feather(\"tags\")\n",
+    "posts_by_tag = get_feather(\"posts_by_tag\").set_index(\"tag_id\")\n",
+    "tags_by_post = get_feather(\"tags_by_post\").set_index(\"post_id\")\n",
+    "implications = get_feather(\"implications\")\n",
     "tags_by_name = tags.copy(deep=True)\n",
     "tags_by_name.set_index(\"name\", inplace=True)\n",
     "tags.set_index(\"tag_id\", inplace=True)"