support specifying which files to download (#33)

substratusai · Aug 23, 2023 · 95f2a91 · 95f2a91
1 parent c4a2f29
commit 95f2a91
Showing 1 changed file with 7 additions and 36 deletions.
diff --git a/model-loader-huggingface/src/load.ipynb b/model-loader-huggingface/src/load.ipynb
@@ -61,45 +61,16 @@
    ],
    "source": [
     "from huggingface_hub.hf_api import model_info\n",
+    "from utils import filter_files\n",
     "\n",
     "model = model_info(model_id)\n",
     "\n",
-    "filenames = [f.rfilename for f in model.siblings ]\n",
-    "filenames"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "7983d9be",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['.gitattributes',\n",
-       " 'README.md',\n",
-       " 'config.json',\n",
-       " 'configuration_RW.py',\n",
-       " 'generation_config.json',\n",
-       " 'handler.py',\n",
-       " 'modelling_RW.py',\n",
-       " 'pytorch_model-00001-of-00002.bin',\n",
-       " 'pytorch_model-00002-of-00002.bin',\n",
-       " 'pytorch_model.bin.index.json',\n",
-       " 'special_tokens_map.json',\n",
-       " 'tokenizer.json',\n",
-       " 'tokenizer_config.json']"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "from utils import filter_files\n",
-    "filenames = filter_files(filenames)\n",
+    "files = os.environ.get(\"PARAM_FILES\", \"\")\n",
+    "if files:\n",
+    "    filenames = [f.strip() for f in files.split(\",\")]\n",
+    "else:\n",
+    "    filenames = [f.rfilename for f in model.siblings ]\n",
+    "    filenames = filter_files(filenames)\n",
     "filenames"
    ]
   },