Add relative path functionality for colab

test_pretrained.ipynb

@@ -35,17 +35,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "current_path = \"./\"\n",
     "\n",
+    "def get_path(rel_path):\n",
+    "    return os.path.join(current_path, rel_path)\n",
+    "\n",
     "if is_google_colab:\n",
     "    hugging_face_path = snapshot_download(\n",
     "        repo_id=\"USC-Applied-NLP-Group/SQL-Generation\",\n",
     "        repo_type=\"model\",  \n",
-    "        allow_patterns=[\"src/*\"],  \n",
+    "        allow_patterns=[\"src/*\", \"train-data/*\", \"deepseek-coder-1.3b-instruct/*\"], \n",
     "    )\n",
     "    sys.path.append(hugging_face_path)\n",
     "    current_path = hugging_face_path"
@@ -70,7 +73,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 36,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -88,7 +91,7 @@
    ],
    "source": [
     "# Load dataset and check length\n",
-    "df = pd.read_csv(os.path.join(current_path, \"train-data/sql_train.tsv\"), sep=\"\\t\")\n",
+    "df = pd.read_csv(get_path(\"train-data/sql_train.tsv\"), sep=\"\\t\")\n",
     "print(\"Total dataset examples: \" + str(len(df)))\n",
     "print(\"\\n\")\n",
     "\n",
@@ -108,7 +111,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -116,8 +119,8 @@
     "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "\n",
     "# Load model and tokenizer\n",
-    "tokenizer = AutoTokenizer.from_pretrained(\"./deepseek-coder-1.3b-instruct\")\n",
-    "model = AutoModelForCausalLM.from_pretrained(\"./deepseek-coder-1.3b-instruct\", torch_dtype=torch.bfloat16, device_map=device) \n",
+    "tokenizer = AutoTokenizer.from_pretrained(get_path(\"deepseek-coder-1.3b-instruct\"))\n",
+    "model = AutoModelForCausalLM.from_pretrained(get_path(\"deepseek-coder-1.3b-instruct\"), torch_dtype=torch.bfloat16, device_map=device) \n",
     "model.generation_config.pad_token_id = tokenizer.pad_token_id"
    ]
   },