|
| 1 | +{ |
| 2 | + "cells": [ |
| 3 | + { |
| 4 | + "cell_type": "code", |
| 5 | + "execution_count": 1, |
| 6 | + "metadata": {}, |
| 7 | + "outputs": [], |
| 8 | + "source": [ |
| 9 | + "import json\n", |
| 10 | + "import sqlparse\n", |
| 11 | + "import pickle as pkl\n", |
| 12 | + "dataset_names = ['academic', 'atis', 'advising', 'geography', 'imdb', 'restaurants', 'scholar', 'yelp']\n", |
| 13 | + "\n", |
| 14 | + "# these datasets are small, so we use the full set. \n", |
| 15 | + "new_split_defined = {'restaurants', 'academic', 'imdb', 'yelp'} " |
| 16 | + ] |
| 17 | + }, |
| 18 | + { |
| 19 | + "cell_type": "code", |
| 20 | + "execution_count": 2, |
| 21 | + "metadata": {}, |
| 22 | + "outputs": [], |
| 23 | + "source": [ |
| 24 | + "# loading the original datasets from the paper:\n", |
| 25 | + "# Improving Text-to-SQL Evaluation Methodology\n", |
| 26 | + "\n", |
| 27 | + "# a dataset is a list of dictionaries\n", |
| 28 | + "# in the original dictionary, each datapoint might consist of several natural language sentences or SQL\n", |
| 29 | + "orig_datasets = []\n", |
| 30 | + "for dataset_name in dataset_names:\n", |
| 31 | + " orig_dataset = json.load(open('text2sql-data/data/%s.json' % dataset_name))\n", |
| 32 | + " for idx, d in enumerate(orig_dataset):\n", |
| 33 | + " \n", |
| 34 | + " d['orig_id'] = (dataset_name, idx)\n", |
| 35 | + " \n", |
| 36 | + " # fixing annotations here\n", |
| 37 | + " \n", |
| 38 | + " # change \"company_name\" to producer name, otherwise there is no variable to replace\n", |
| 39 | + " if dataset_name == 'imdb' and idx == 27:\n", |
| 40 | + " d['sql'][0] = 'SELECT MOVIEalias0.TITLE FROM COMPANY AS COMPANYalias0 , COPYRIGHT AS COPYRIGHTalias0 , MOVIE AS MOVIEalias0 WHERE COMPANYalias0.NAME = \"producer_name0\" AND COPYRIGHTalias0.CID = COMPANYalias0.ID AND MOVIEalias0.MID = COPYRIGHTalias0.MSID AND MOVIEalias0.RELEASE_YEAR > movie_release_year0 ;'\n", |
| 41 | + " \n", |
| 42 | + " # removing the extra space surrounding the variable actor_name0\n", |
| 43 | + " if dataset_name == 'imdb' and idx == 78:\n", |
| 44 | + " d['sql'][0] = 'SELECT MAX( DERIVED_TABLEalias0.DERIVED_FIELDalias0 ) FROM ( SELECT COUNT( DISTINCT ( MOVIEalias0.TITLE ) ) AS DERIVED_FIELDalias0 FROM ACTOR AS ACTORalias0 , CAST AS CASTalias0 , MOVIE AS MOVIEalias0 WHERE ACTORalias0.NAME = \"actor_name0\" AND CASTalias0.AID = ACTORalias0.AID AND MOVIEalias0.MID = CASTalias0.MSID GROUP BY MOVIEalias0.RELEASE_YEAR ) AS DERIVED_TABLEalias0 ;'\n", |
| 45 | + " \n", |
| 46 | + " # there was a scoping error; changed AUTHORalias1 to AUTHORalias0, PUBLICATIONalias1 to PUBLICATIONalias0\n", |
| 47 | + " if dataset_name == 'academic' and idx == 182:\n", |
| 48 | + " d['sql'][0] = 'SELECT DERIVED_FIELDalias0 FROM ( SELECT AUTHORalias0.NAME AS DERIVED_FIELDalias0 , COUNT( DISTINCT ( PUBLICATIONalias0.TITLE ) ) AS DERIVED_FIELDalias1 FROM AUTHOR AS AUTHORalias0 , CONFERENCE AS CONFERENCEalias0 , PUBLICATION AS PUBLICATIONalias0 , WRITES AS WRITESalias0 WHERE CONFERENCEalias0.NAME = \"conference_name0\" AND PUBLICATIONalias0.CID = CONFERENCEalias0.CID AND WRITESalias0.AID = AUTHORalias0.AID AND WRITESalias0.PID = PUBLICATIONalias0.PID GROUP BY AUTHORalias0.NAME ) AS DERIVED_TABLEalias0 , ( SELECT AUTHORalias1.NAME AS DERIVED_FIELDalias2 , COUNT( DISTINCT ( PUBLICATIONalias1.TITLE ) ) AS DERIVED_FIELDalias3 FROM AUTHOR AS AUTHORalias1 , CONFERENCE AS CONFERENCEalias1 , PUBLICATION AS PUBLICATIONalias1 , WRITES AS WRITESalias1 WHERE CONFERENCEalias1.NAME = \"conference_name1\" AND PUBLICATIONalias1.CID = CONFERENCEalias1.CID AND WRITESalias1.AID = AUTHORalias1.AID AND WRITESalias1.PID = PUBLICATIONalias1.PID GROUP BY AUTHORalias1.NAME ) AS DERIVED_TABLEalias1 WHERE DERIVED_TABLEalias0.DERIVED_FIELDalias1 > DERIVED_TABLEalias1.DERIVED_FIELDalias3 AND DERIVED_TABLEalias1.DERIVED_FIELDalias2 = DERIVED_TABLEalias0.DERIVED_FIELDalias0 ;'\n", |
| 49 | + " \n", |
| 50 | + " # wrong number of arguments to function COUNT(), change from \",\" to \"||\" for sqlite3 to recognize and execute\n", |
| 51 | + " if dataset_name == 'advising' and idx == 107:\n", |
| 52 | + " d['sql'][0] = 'SELECT COUNT( DISTINCT COURSEalias1.DEPARTMENT || COURSEalias0.NUMBER ) FROM COURSE AS COURSEalias0 , COURSE AS COURSEalias1 , COURSE_PREREQUISITE AS COURSE_PREREQUISITEalias0 , STUDENT_RECORD AS STUDENT_RECORDalias0 WHERE COURSEalias0.COURSE_ID = COURSE_PREREQUISITEalias0.PRE_COURSE_ID AND COURSEalias1.COURSE_ID = COURSE_PREREQUISITEalias0.COURSE_ID AND COURSEalias1.DEPARTMENT = \"department0\" AND COURSEalias1.NUMBER = number0 AND STUDENT_RECORDalias0.COURSE_ID = COURSEalias0.COURSE_ID AND STUDENT_RECORDalias0.STUDENT_ID = 1 ;'\n", |
| 53 | + " \n", |
| 54 | + " # there was not example given for level1 and hence replacing variable with values leads to errors\n", |
| 55 | + " if dataset_name == 'advising' and idx == 132:\n", |
| 56 | + " d['variables'][0]['example'] = '300'\n", |
| 57 | + " \n", |
| 58 | + " # cannot use count and order without group by; added grouping by actor_id\n", |
| 59 | + " if dataset_name == 'imdb' and idx == 79:\n", |
| 60 | + " d['sql'][0] = 'SELECT ACTORalias0.NAME FROM ACTOR AS ACTORalias0 , CAST AS CASTalias0 , MOVIE AS MOVIEalias0 WHERE CASTalias0.AID = ACTORalias0.AID AND MOVIEalias0.MID = CASTalias0.MSID GROUP BY ACTORalias0.AID ORDER BY COUNT( DISTINCT ( MOVIEalias0.TITLE ) ) DESC LIMIT 1 ;'\n", |
| 61 | + " \n", |
| 62 | + " # cannot use count and order without group by; added grouping by actor_id\n", |
| 63 | + " if dataset_name == 'imdb' and idx == 80:\n", |
| 64 | + " d['sql'][0] = 'SELECT ACTORalias0.NAME FROM ACTOR AS ACTORalias0 , CAST AS CASTalias0 , DIRECTED_BY AS DIRECTED_BYalias0 , DIRECTOR AS DIRECTORalias0 , MOVIE AS MOVIEalias0 WHERE CASTalias0.AID = ACTORalias0.AID AND DIRECTORalias0.DID = DIRECTED_BYalias0.DID AND MOVIEalias0.MID = CASTalias0.MSID AND MOVIEalias0.MID = DIRECTED_BYalias0.MSID GROUP BY ACTORalias0.AID ORDER BY COUNT( DISTINCT ( MOVIEalias0.TITLE ) ) DESC LIMIT 1 ;'\n", |
| 65 | + " \n", |
| 66 | + " # table has \"u\" in the neighborhood spelling.\n", |
| 67 | + " n_before, n_after = 'NEIGHBORHOOD', 'NEIGHBOURHOOD'\n", |
| 68 | + " if dataset_name == 'yelp':\n", |
| 69 | + " d['sql'][0] = d['sql'][0].replace(n_before, n_after)\n", |
| 70 | + " \n", |
| 71 | + " if dataset_name == 'yelp' and idx == 42:\n", |
| 72 | + " d['sql'][0] = 'SELECT NEIGHBOURHOODalias0.NEIGHBOURHOOD_NAME FROM BUSINESS AS BUSINESSalias0 , NEIGHBOURHOOD AS NEIGHBOURHOODalias0 , REVIEW AS REVIEWalias0 , USER AS USERalias0 WHERE NEIGHBOURHOODalias0.BUSINESS_ID = BUSINESSalias0.BUSINESS_ID AND REVIEWalias0.BUSINESS_ID = BUSINESSalias0.BUSINESS_ID AND USERalias0.NAME = \"user_name0\" AND USERalias0.USER_ID = REVIEWalias0.USER_ID ;'\n", |
| 73 | + "\n", |
| 74 | + " orig_datasets.extend(orig_dataset)" |
| 75 | + ] |
| 76 | + }, |
| 77 | + { |
| 78 | + "cell_type": "code", |
| 79 | + "execution_count": 3, |
| 80 | + "metadata": {}, |
| 81 | + "outputs": [ |
| 82 | + { |
| 83 | + "name": "stdout", |
| 84 | + "output_type": "stream", |
| 85 | + "text": [ |
| 86 | + "There are 3509 datapoints in the new testset\n" |
| 87 | + ] |
| 88 | + } |
| 89 | + ], |
| 90 | + "source": [ |
| 91 | + "# we create the new testset here\n", |
| 92 | + "new_testset = []\n", |
| 93 | + "for d in orig_datasets:\n", |
| 94 | + " orig_id = d['orig_id']\n", |
| 95 | + " db_id, idx = orig_id\n", |
| 96 | + " \n", |
| 97 | + " # we only incorporate the test split if the dataset is large enough\n", |
| 98 | + " # otherwise we incorporate the entire dataset\n", |
| 99 | + " if d['query-split'] != 'test' and db_id not in new_split_defined:\n", |
| 100 | + " continue\n", |
| 101 | + " sql = d['sql'][0]\n", |
| 102 | + " instance_variables = d['variables']\n", |
| 103 | + " instance_name2examples = {d['name']: d['example'] for d in instance_variables}\n", |
| 104 | + " \n", |
| 105 | + " # we create a new datapoint for each natural language query\n", |
| 106 | + " for sentence in d['sentences']:\n", |
| 107 | + " new_datapoint = {\n", |
| 108 | + " 'text': sentence['text'],\n", |
| 109 | + " 'query': sql,\n", |
| 110 | + " 'variables': instance_variables,\n", |
| 111 | + " 'orig_id': orig_id,\n", |
| 112 | + " 'db_id': db_id,\n", |
| 113 | + " 'db_path': 'database/{db_id}/{db_id}.sqlite'.format(db_id=db_id)\n", |
| 114 | + " }\n", |
| 115 | + " new_testset.append(new_datapoint)\n", |
| 116 | + "print('There are %d datapoints in the new testset' % len(new_testset))" |
| 117 | + ] |
| 118 | + }, |
| 119 | + { |
| 120 | + "cell_type": "code", |
| 121 | + "execution_count": 4, |
| 122 | + "metadata": {}, |
| 123 | + "outputs": [], |
| 124 | + "source": [ |
| 125 | + "import re\n", |
| 126 | + "\n", |
| 127 | + "# this block implements a function that extract variable names from text and sql\n", |
| 128 | + "# later we use it to ensure that every variable is replaced\n", |
| 129 | + "\n", |
| 130 | + "variable_pattern = re.compile('^[a-z_]+[0-9]+$')\n", |
| 131 | + "\n", |
| 132 | + "def extract_variable_names(t):\n", |
| 133 | + " tokens = t.replace('\"', '').replace('%', '').split(' ')\n", |
| 134 | + " var_names = {v for v in tokens if variable_pattern.match(v) and 'alias' not in v}\n", |
| 135 | + " return var_names\n", |
| 136 | + "\n", |
| 137 | + "test = False\n", |
| 138 | + "if test:\n", |
| 139 | + " sql = 'SELECT BUSINESSalias0.NAME FROM BUSINESS AS BUSINESSalias0 , REVIEW AS REVIEWalias0 WHERE REVIEWalias0.BUSINESS_ID = BUSINESSalias0.BUSINESS_ID AND REVIEWalias0.MONTH = \"review_month0\" GROUP BY BUSINESSalias0.NAME ORDER BY COUNT( DISTINCT ( REVIEWalias0.TEXT ) ) DESC LIMIT 1 ;'\n", |
| 140 | + " print(extract_variable_names(sql))\n", |
| 141 | + " text = 'return me the homepage of journal_name0 .'\n", |
| 142 | + " print(extract_variable_names(text))" |
| 143 | + ] |
| 144 | + }, |
| 145 | + { |
| 146 | + "cell_type": "code", |
| 147 | + "execution_count": 5, |
| 148 | + "metadata": {}, |
| 149 | + "outputs": [], |
| 150 | + "source": [ |
| 151 | + "# this block removes extra space surrounding variable names\n", |
| 152 | + "def remove_extra_space_around_variable(t):\n", |
| 153 | + " var_names = extract_variable_names(t)\n", |
| 154 | + " result = str(t)\n", |
| 155 | + " for v in var_names:\n", |
| 156 | + " result = result.replace('\" ' + v + ' \"', v)\n", |
| 157 | + " return result" |
| 158 | + ] |
| 159 | + }, |
| 160 | + { |
| 161 | + "cell_type": "code", |
| 162 | + "execution_count": 6, |
| 163 | + "metadata": {}, |
| 164 | + "outputs": [ |
| 165 | + { |
| 166 | + "name": "stdout", |
| 167 | + "output_type": "stream", |
| 168 | + "text": [ |
| 169 | + "set()\n" |
| 170 | + ] |
| 171 | + } |
| 172 | + ], |
| 173 | + "source": [ |
| 174 | + "problematic = set()\n", |
| 175 | + "\n", |
| 176 | + "for datapoint in new_testset:\n", |
| 177 | + " orig_id = datapoint['orig_id']\n", |
| 178 | + " \n", |
| 179 | + " # remove extra whitespace surrounding the text\n", |
| 180 | + " datapoint['text'] = remove_extra_space_around_variable(datapoint['text'])\n", |
| 181 | + " \n", |
| 182 | + " # there should not be extra whitespace surrounding the sql variables\n", |
| 183 | + " if datapoint['query'] != remove_extra_space_around_variable(datapoint['query']):\n", |
| 184 | + " problematic.add(orig_id)\n", |
| 185 | + "\n", |
| 186 | + " text_vars = extract_variable_names(datapoint['text'])\n", |
| 187 | + " sql_vars = extract_variable_names(datapoint['query'])\n", |
| 188 | + " \n", |
| 189 | + " instance_variables = {d['name']: d for d in datapoint['variables']}\n", |
| 190 | + " \n", |
| 191 | + " # we ensure that all the variables in the sql query and the text can be replaced\n", |
| 192 | + " # by some variable in the variable dictionary\n", |
| 193 | + " if len(text_vars - instance_variables.keys()) != 0 or len(sql_vars - instance_variables.keys()):\n", |
| 194 | + " problematic.add(orig_id)\n", |
| 195 | + " \n", |
| 196 | + " # replace the variables with the examples in the variable dictionary\n", |
| 197 | + " for text_var in text_vars:\n", |
| 198 | + " datapoint['text'] = datapoint['text'].replace(text_var, instance_variables[text_var]['example'])\n", |
| 199 | + " \n", |
| 200 | + " for sql_var in sql_vars:\n", |
| 201 | + " datapoint['query'] = datapoint['query'].replace(sql_var, instance_variables[sql_var]['example'])\n", |
| 202 | + "\n", |
| 203 | + "# we can trace back which datapoints do not satisfy the assumption,\n", |
| 204 | + "# then go back and fix it manually\n", |
| 205 | + "print(problematic)" |
| 206 | + ] |
| 207 | + }, |
| 208 | + { |
| 209 | + "cell_type": "code", |
| 210 | + "execution_count": 7, |
| 211 | + "metadata": {}, |
| 212 | + "outputs": [ |
| 213 | + { |
| 214 | + "name": "stdout", |
| 215 | + "output_type": "stream", |
| 216 | + "text": [ |
| 217 | + "[{'db_id': 'academic',\n", |
| 218 | + " 'db_path': 'database/academic/academic.sqlite',\n", |
| 219 | + " 'orig_id': ('academic', 0),\n", |
| 220 | + " 'query': 'SELECT JOURNALalias0.HOMEPAGE FROM JOURNAL AS JOURNALalias0 WHERE '\n", |
| 221 | + " 'JOURNALalias0.NAME = \"PVLDB\" ;',\n", |
| 222 | + " 'text': 'return me the homepage of PVLDB .',\n", |
| 223 | + " 'variables': [{'example': 'PVLDB',\n", |
| 224 | + " 'location': 'both',\n", |
| 225 | + " 'name': 'journal_name0',\n", |
| 226 | + " 'type': 'journal_name'}]},\n", |
| 227 | + " {'db_id': 'academic',\n", |
| 228 | + " 'db_path': 'database/academic/academic.sqlite',\n", |
| 229 | + " 'orig_id': ('academic', 1),\n", |
| 230 | + " 'query': 'SELECT AUTHORalias0.HOMEPAGE FROM AUTHOR AS AUTHORalias0 WHERE '\n", |
| 231 | + " 'AUTHORalias0.NAME = \"H. V. Jagadish\" ;',\n", |
| 232 | + " 'text': 'return me the homepage of H. V. Jagadish .',\n", |
| 233 | + " 'variables': [{'example': 'H. V. Jagadish',\n", |
| 234 | + " 'location': 'both',\n", |
| 235 | + " 'name': 'author_name0',\n", |
| 236 | + " 'type': 'author_name'}]},\n", |
| 237 | + " {'db_id': 'academic',\n", |
| 238 | + " 'db_path': 'database/academic/academic.sqlite',\n", |
| 239 | + " 'orig_id': ('academic', 2),\n", |
| 240 | + " 'query': 'SELECT PUBLICATIONalias0.ABSTRACT FROM PUBLICATION AS '\n", |
| 241 | + " 'PUBLICATIONalias0 WHERE PUBLICATIONalias0.TITLE = \"Making database '\n", |
| 242 | + " 'systems usable\" ;',\n", |
| 243 | + " 'text': 'return me the abstract of Making database systems usable .',\n", |
| 244 | + " 'variables': [{'example': 'Making database systems usable',\n", |
| 245 | + " 'location': 'both',\n", |
| 246 | + " 'name': 'publication_title0',\n", |
| 247 | + " 'type': 'publication_title'}]},\n", |
| 248 | + " {'db_id': 'academic',\n", |
| 249 | + " 'db_path': 'database/academic/academic.sqlite',\n", |
| 250 | + " 'orig_id': ('academic', 3),\n", |
| 251 | + " 'query': 'SELECT PUBLICATIONalias0.YEAR FROM PUBLICATION AS '\n", |
| 252 | + " 'PUBLICATIONalias0 WHERE PUBLICATIONalias0.TITLE = \"Making database '\n", |
| 253 | + " 'systems usable\" ;',\n", |
| 254 | + " 'text': 'return me the year of Making database systems usable',\n", |
| 255 | + " 'variables': [{'example': 'Making database systems usable',\n", |
| 256 | + " 'location': 'both',\n", |
| 257 | + " 'name': 'publication_title0',\n", |
| 258 | + " 'type': 'publication_title'}]},\n", |
| 259 | + " {'db_id': 'academic',\n", |
| 260 | + " 'db_path': 'database/academic/academic.sqlite',\n", |
| 261 | + " 'orig_id': ('academic', 3),\n", |
| 262 | + " 'query': 'SELECT PUBLICATIONalias0.YEAR FROM PUBLICATION AS '\n", |
| 263 | + " 'PUBLICATIONalias0 WHERE PUBLICATIONalias0.TITLE = \"Making database '\n", |
| 264 | + " 'systems usable\" ;',\n", |
| 265 | + " 'text': 'return me the year of Making database systems usable .',\n", |
| 266 | + " 'variables': [{'example': 'Making database systems usable',\n", |
| 267 | + " 'location': 'both',\n", |
| 268 | + " 'name': 'publication_title0',\n", |
| 269 | + " 'type': 'publication_title'}]}]\n" |
| 270 | + ] |
| 271 | + } |
| 272 | + ], |
| 273 | + "source": [ |
| 274 | + "from pprint import pprint\n", |
| 275 | + "\n", |
| 276 | + "pprint(new_testset[:5])" |
| 277 | + ] |
| 278 | + } |
| 279 | + ], |
| 280 | + "metadata": { |
| 281 | + "kernelspec": { |
| 282 | + "display_name": "Python 3", |
| 283 | + "language": "python", |
| 284 | + "name": "python3" |
| 285 | + }, |
| 286 | + "language_info": { |
| 287 | + "codemirror_mode": { |
| 288 | + "name": "ipython", |
| 289 | + "version": 3 |
| 290 | + }, |
| 291 | + "file_extension": ".py", |
| 292 | + "mimetype": "text/x-python", |
| 293 | + "name": "python", |
| 294 | + "nbconvert_exporter": "python", |
| 295 | + "pygments_lexer": "ipython3", |
| 296 | + "version": "3.7.4" |
| 297 | + } |
| 298 | + }, |
| 299 | + "nbformat": 4, |
| 300 | + "nbformat_minor": 2 |
| 301 | +} |
0 commit comments