Add support for colab

test_pretrained.ipynb

@@ -7,6 +7,56 @@
     "# Run pre-trained DeepSeek Coder 1.3B Model on Chat-GPT 4o generated dataset"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd \n",
+    "import warnings\n",
+    "warnings.filterwarnings(\"ignore\")\n",
+    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+    "import torch\n",
+    "import sys\n",
+    "import sqlite3 as sql\n",
+    "from huggingface_hub import snapshot_download"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "is_google_colab=False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "    )\n",
+    "    sys.path.append(hugging_face_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from src.prompts.prompt import input_text\n",
+    "from src.evaluation.compare_result import compare_result"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -33,10 +83,6 @@
     }
    ],
    "source": [
-    "import pandas as pd \n",
-    "import warnings\n",
-    "warnings.filterwarnings(\"ignore\")\n",
-    "\n",
     "# Load dataset and check length\n",
     "df = pd.read_csv(\"./train-data/sql_train.tsv\", sep='\\t')\n",
     "print(\"Total dataset examples: \" + str(len(df)))\n",
@@ -62,9 +108,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
-    "import torch\n",
-    "\n",
     "# Set device to cuda if available, otherwise CPU\n",
     "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "\n",
@@ -74,22 +117,6 @@
     "model.generation_config.pad_token_id = tokenizer.pad_token_id"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Create prompt to setup the model for better performance"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from src.prompts.prompt import input_text"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -144,8 +171,6 @@
     }
    ],
    "source": [
-    "import sqlite3 as sql\n",
-    "\n",
     "# Create connection to sqlite3 database\n",
     "connection = sql.connect('./nba-data/nba.sqlite')\n",
     "cursor = connection.cursor()\n",
@@ -193,115 +218,12 @@
     }
    ],
    "source": [
-    "import math\n",
-    "from src.evaluation.compare_result import compare_result_two\n",
-    "\n",
-    "def compare_result(sample_query, sample_result, query_output):\n",
-    "    # Clean model output to only have the query output\n",
-    "    if query_output[0:7] == \"SQLite:\":\n",
-    "        query = query_output[7:]\n",
-    "    elif query_output[0:4] == \"SQL:\":\n",
-    "        query = query_output[4:]\n",
-    "    else:\n",
-    "        query = query_output\n",
-    "    \n",
-    "    # Try to execute query, if it fails, then this is a failure of the model\n",
-    "    try:\n",
-    "        # Execute query and obtain result\n",
-    "        cursor.execute(query)\n",
-    "        rows = cursor.fetchall()\n",
-    "\n",
-    "        # Strip all whitespace before comparing queries since there may be differences in spacing, newlines, tabs, etc.\n",
-    "        query = query.replace(\" \", \"\").replace(\"\\n\", \"\").replace(\"\\t\", \"\")\n",
-    "        sample_query = sample_query.replace(\" \", \"\").replace(\"\\n\", \"\").replace(\"\\t\", \"\")\n",
-    "        query_match = (query == sample_query)\n",
-    "\n",
-    "        # If the queries match, the results clearly also match\n",
-    "        if query_match:\n",
-    "            return True, True, True\n",
-    "\n",
-    "        # Check if this is a multi-line query\n",
-    "        if \"|\" in sample_result or \"(\" in sample_result:\n",
-    "            #print(rows)\n",
-    "            # Create list of results by stripping separators and splitting on them\n",
-    "            if \"(\" in sample_result:\n",
-    "                sample_result = sample_result.replace(\"(\", \"\").replace(\")\", \"\")\n",
-    "                result_list = sample_result.split(\",\") \n",
-    "            else:\n",
-    "                result_list = sample_result.split(\"|\") \n",
-    "\n",
-    "            # Strip all results in list\n",
-    "            for i in range(len(result_list)):\n",
-    "                result_list[i] = str(result_list[i]).strip()\n",
-    "            \n",
-    "            # Loop through model result and see if it matches training example\n",
-    "            result = False\n",
-    "            for row in rows:\n",
-    "                for r in row:\n",
-    "                    for res in result_list:\n",
-    "                        try:\n",
-    "                            if math.isclose(float(r), float(res), abs_tol=0.5):\n",
-    "                                return True, query_match, True\n",
-    "                        except:\n",
-    "                            if r in res or res in r:\n",
-    "                                return True, query_match, True\n",
-    "                    \n",
-    "            # Check if the model returned a sum of examples as opposed to the whole thing\n",
-    "            if len(rows) == 1:\n",
-    "                for r in rows[0]:\n",
-    "                    if r == str(len(result_list)):\n",
-    "                        return True, query_match, True\n",
-    "                    \n",
-    "            return True, query_match, result\n",
-    "        # Else the sample result is a single value or string\n",
-    "        else:\n",
-    "            #print(rows)\n",
-    "            result = False\n",
-    "            # Loop through model result and see if it contains the sample result\n",
-    "            for row in rows:\n",
-    "                for r in row:\n",
-    "                    # Check by string\n",
-    "                    if str(r) in str(sample_result):\n",
-    "                        try:\n",
-    "                            if math.isclose(float(r), float(sample_result), abs_tol=0.5):\n",
-    "                                return True, query_match, True\n",
-    "                        except:\n",
-    "                            return True, query_match, True\n",
-    "                    # Check by number, using try incase the cast as float fails\n",
-    "                    try:\n",
-    "                        if math.isclose(float(r), float(sample_result), abs_tol=0.5):\n",
-    "                            return True, query_match, True\n",
-    "                    except:\n",
-    "                        pass\n",
-    "\n",
-    "            # Check if the model returned a list of examples instead of a total sum (both acceptable)\n",
-    "            try:\n",
-    "                if len(rows) > 1 and len(rows) == int(sample_result):\n",
-    "                    return True, query_match, True\n",
-    "                if len(rows[0]) > 1 and rows[0][1] is not None and  len(rows[0]) == int(sample_result):\n",
-    "                    return True, query_match, True\n",
-    "            except:\n",
-    "                pass\n",
-    "\n",
-    "            # Compare results and return\n",
-    "            return True, query_match, result\n",
-    "    except:\n",
-    "        return False, False, False\n",
-    "\n",
     "# Obtain sample\n",
     "sample = df.sample(n=1)\n",
-    "sample_dic = {\n",
-    "    \"natural_query\": \"How many home games did the Miami Heat play in the 2021 season?\",\n",
-    "    \"sql_query\": \"SELECT COUNT(*) FROM game WHERE team_name_home = 'Miami Heat' AND season_id = '22021';\",\n",
-    "    \"result\": 41.0\n",
-    "}\n",
     "\n",
-    "sample = pd.DataFrame([sample_dic])\n",
-    "\"\"\"\n",
     "print(sample[\"natural_query\"].values[0])\n",
     "print(sample[\"sql_query\"].values[0])\n",
     "print(sample[\"result\"].values[0])\n",
-    "\"\"\"\n",
     "\n",
     "# Create message with sample query and run model\n",
     "message=[{ 'role': 'user', 'content': input_text + sample[\"natural_query\"].values[0]}]\n",
@@ -312,15 +234,10 @@
     "query_output = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)\n",
     "print(query_output)\n",
     "\n",
-    "result = compare_result(sample[\"sql_query\"].values[0], sample[\"result\"].values[0], query_output)\n",
+    "result = compare_result(cursor, sample[\"sql_query\"].values[0], sample[\"result\"].values[0], query_output)\n",
     "print(\"Statement valid? \" + str(result[0]))\n",
     "print(\"SQLite matched? \" + str(result[1]))\n",
-    "print(\"Result matched? \" + str(result[2]))\n",
-    "\n",
-    "result_two = compare_result_two(cursor, sample[\"sql_query\"].values[0], sample[\"result\"].values[0], query_output)\n",
-    "print(\"Statement valid? \" + str(result_two[0]))\n",
-    "print(\"SQLite matched? \" + str(result_two[1]))\n",
-    "print(\"Result matched? \" + str(result_two[2]))"
+    "print(\"Result matched? \" + str(result[2]))"
    ]
   },
   {