kaggle-bert-nl2sql.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":7335893,"sourceType":"datasetVersion","datasetId":4258769},{"sourceId":7335929,"sourceType":"datasetVersion","datasetId":4258788},{"sourceId":7359762,"sourceType":"datasetVersion","datasetId":4274874},{"sourceId":7462076,"sourceType":"datasetVersion","datasetId":4343089},{"sourceId":7467421,"sourceType":"datasetVersion","datasetId":4346710}],"dockerImageVersionId":30627,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#import\nimport numpy as np\nimport pandas as pd\nimport torch.utils.data\nimport torch.cuda\nimport json\nimport re\nimport jieba.posseg as psg\nimport nltk\nfrom transformers import BertTokenizer,BertModel\nfrom torch import nn\nfrom sklearn import metrics\nfrom torch.optim.adam import Adam\nfrom torch.utils.data import DataLoader, random_split\nfrom tqdm import tqdm\nfrom typing import List\nfrom datetime import datetime, timedelta\nfrom functools import reduce\nfrom dateutil.relativedelta import relativedelta\nfrom nltk.corpus import stopwords","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-01-24T03:23:30.647834Z","iopub.execute_input":"2024-01-24T03:23:30.648205Z","iopub.status.idle":"2024-01-24T03:23:35.003407Z","shell.execute_reply.started":"2024-01-24T03:23:30.648171Z","shell.execute_reply":"2024-01-24T03:23:35.002580Z"},"trusted":true},"execution_count":1,"outputs":[{"name":"stderr","text":"/opt/conda/lib/python3.10/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.24.3\n  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n","output_type":"stream"}]},{"cell_type":"code","source":"#time_extract\ndef time_extract(text):\n    replace_text = text\n    time_res = []\n    word = ''\n    key_date = {'前天': -2, '前日': -2, '昨天': -2, '昨日': -2, '今天': 0, '今日': 0, '明天': 1, '明日': 1, '后天': 2,\n                '后日': 2}\n    key_month = {'上月': -1, '上个月': -1, '这月': 0, '本月': 0, '下月': 1, '下个月': 1}\n    key_year = {'去年': -1, '今年': 0, '本年': 0, '明年': 1}\n\n    for k, v in psg.cut(text):\n        if k in key_date:\n            if word != '':\n                time_res.append(word)\n            offset_date = (datetime.today() + timedelta(days=key_date.get(k, 0)))\n            replace_text = text.replace(k, offset_date.strftime('%Y-%m-%d'))\n            word = offset_date.strftime('%Y年%m月%d日')\n        elif k in key_month:\n            if word != '':\n                time_res.append(word)\n            offset_date = (datetime.today() + relativedelta(months=key_month.get(k, 0)))\n            replace_text = text.replace(k, offset_date.strftime('%Y-%m'))\n            word = offset_date.strftime('%Y年%m月')\n        elif k in key_year:\n            if word != '':\n                time_res.append(word)\n            offset_date = (datetime.today() + relativedelta(years=key_year.get(k, 0)))\n            replace_text = text.replace(k, offset_date.strftime('%Y'))\n            word = offset_date.strftime('%Y年')\n        elif word != '':\n            if v in ['m', 't']:\n                word = word + k\n            else:\n                time_res.append(word)\n                word = ''\n        elif v in ['m', 't']:\n            word = k\n    if word != '':\n        time_res.append(word)\n    result = list(filter(lambda x: x is not None, [check_time_valid(w) for w in time_res]))\n    final_res = [parse_datetime(w) for w in result]\n    replace_text = reduce(lambda text, replacement: text.replace(*replacement) if replacement[1] is not None else text,\n                          zip(time_res, final_res), replace_text)\n    return replace_text, [x for x in final_res if x is not None]\n\n\ndef check_time_valid(word):\n    m = re.match(\"\\d+$\", word)\n    if m:\n        if len(word) <= 6:\n            return None\n    word1 = re.sub('[号|日]\\d+$', '日', word)\n    if word1 != word:\n        return check_time_valid(word1)\n    else:\n        return word1\n\n\ndef parse_datetime(msg):\n    if msg is None or len(msg) == 0:\n        return None\n\n    m = re.match(\n        r\"([0-9零一二两三四五六七八九十]+年)?([0-9一二两三四五六七八九十]+月)?([0-9一二两三四五六七八九十]+[号日])?([上中下午晚早]+)?([0-9零一二两三四五六七八九十百]+[点:\\.时])?([0-9零一二三四五六七八九十百]+分?)?([0-9零一二三四五六七八九十百]+秒)?\",\n        msg)\n    if m.group(0) is not None and m.group(0) is not '':\n        res = {\n            \"year\": m.group(1),\n            \"month\": m.group(2),\n            \"day\": m.group(3),\n            \"hour\": m.group(5) if m.group(5) is not None else '00',\n            \"minute\": m.group(6) if m.group(6) is not None else '00',\n            \"second\": m.group(7) if m.group(7) is not None else '00',\n        }\n        params = {}\n\n        for name in res:\n            if res[name] is not None and len(res[name]) != 0:\n                if name == 'year':\n                    tmp = year2dig(res[name][:-1])\n                else:\n                    tmp = cn2dig(res[name][:-1])\n                if tmp is not None:\n                    params[name] = int(tmp)\n        target_date = datetime.today().replace(**params)\n        is_pm = m.group(4)\n        if is_pm is not None:\n            if is_pm == u'下午' or is_pm == u'晚上' or is_pm == '中午':\n                hour = target_date.time().hour\n                if hour < 12:\n                    target_date = target_date.replace(hour=hour + 12)\n        formatted_result = None\n        if res['day'] is None and res['month'] is None:\n            formatted_result = target_date.strftime('%Y')\n        elif res['day'] is None and res['month'] is not None:\n            formatted_result = target_date.strftime('%Y-%m')\n        elif res['day'] is not None and res['hour'] == '00' and res['minute'] == '00' and res['second'] == '00':\n            formatted_result = target_date.strftime('%Y-%m-%d')\n        elif res['day'] is not None:\n            formatted_result = target_date.strftime('%Y-%m-%d %H:%M:%S')\n        return formatted_result\n    else:\n        return None\n\n\nUTIL_CN_NUM = {\n    '零': 0, '一': 1, '二': 2, '两': 2, '三': 3, '四': 4,\n    '五': 5, '六': 6, '七': 7, '八': 8, '九': 9,\n    '0': 0, '1': 1, '2': 2, '3': 3, '4': 4,\n    '5': 5, '6': 6, '7': 7, '8': 8, '9': 9\n}\nUTIL_CN_UNIT = {'十': 10, '百': 100, '千': 1000, '万': 10000}\n\n\ndef cn2dig(src):\n    if src == \"\":\n        return None\n    m = re.match(\"\\d+\", src)\n    if m:\n        return int(m.group(0))\n    rsl = 0\n    unit = 1\n    for item in src[::-1]:\n        if item in UTIL_CN_UNIT.keys():\n            unit = UTIL_CN_UNIT[item]\n        elif item in UTIL_CN_NUM.keys():\n            num = UTIL_CN_NUM[item]\n            rsl += num * unit\n        else:\n            return None\n    if rsl < unit:\n        rsl += unit\n    return rsl\n\n\ndef year2dig(year):\n    res = ''\n    for item in year:\n        if item in UTIL_CN_NUM.keys():\n            res = res + str(UTIL_CN_NUM[item])\n        else:\n            res = res + item\n    m = re.match(\"\\d+\", res)\n    if m:\n        if len(m.group(0)) == 2:\n            return int(datetime.today().year / 100) * 100 + int(m.group(0))\n        else:\n            return int(m.group(0))\n    else:\n        return None","metadata":{"execution":{"iopub.status.busy":"2024-01-24T03:23:36.892401Z","iopub.execute_input":"2024-01-24T03:23:36.893209Z","iopub.status.idle":"2024-01-24T03:23:36.924563Z","shell.execute_reply.started":"2024-01-24T03:23:36.893176Z","shell.execute_reply":"2024-01-24T03:23:36.923579Z"},"trusted":true},"execution_count":2,"outputs":[{"name":"stderr","text":"<>:66: SyntaxWarning: \"is not\" with a literal. Did you mean \"!=\"?\n<>:66: SyntaxWarning: \"is not\" with a literal. Did you mean \"!=\"?\n/tmp/ipykernel_130/383107898.py:66: SyntaxWarning: \"is not\" with a literal. Did you mean \"!=\"?\n  if m.group(0) is not None and m.group(0) is not '':\n","output_type":"stream"}]},{"cell_type":"code","source":"#utils\ndef read_train_datas(path, question_length, columns):\n    \"\"\"\n    :param path 数据路径\n    :param question_length 问题长度\n    :param columns 列\n    :return: [[question, agg, conn_op, cond_ops, cond_vals],...], cond_vals:[[val_start_idx,val_end_idx],...]\n    \"\"\"\n    column_length = len(columns)\n    with open(path, 'r', encoding='utf-8') as f:\n        data_list = []\n        for line in f:\n            item = json.loads(line)\n            time_replace(item)\n            # question\n            question = item['question']\n            # agg\n            sel = item['sql']['sel']\n            agg_op = item['sql']['agg']\n            agg = [get_agg_dict()['none']] * column_length\n            for i in range(len(sel)):\n                sel_col_item = sel[i]\n                agg_op_item = agg_op[i]\n                agg[sel_col_item] = agg_op_item\n            # conn_op\n            conn_op = item['sql']['cond_conn_op']\n            # cond_cols & cond_ops & cond_vals\n            # +1 默认初始化为不存在的列, question_length需要大于column_length\n            cond_cols = [column_length + 1] * question_length\n            cond_ops = [get_cond_op_dict()['none']] * question_length\n            cond_vals = [0] * question_length\n            if item['sql'].get('conds') is not None:\n                conds = item['sql']['conds']\n                for idx, cond in enumerate(conds):\n                    cond_cols[idx] = cond[0]\n                    cond_ops[idx] = cond[1]\n                    value = cond[2]\n                    cond_vals = fill_value_start_end(cond_vals, question, value, idx)\n            data_list.append([question, agg, conn_op, cond_cols, cond_ops, cond_vals])\n    return data_list\n\n\ndef read_predict_datas(path):\n    \"\"\"\n    :param path: 预测数据路径\n    :return: 预测数据\n    \"\"\"\n    origin_questions = []\n    questions = []\n    with open(path, 'r', encoding='utf-8') as f:\n        for line in f:\n            item = json.loads(line)\n            origin_question = item['question']\n            origin_questions.append([origin_question])\n            time_replace(item)\n            question = item['question']\n            questions.append([question])\n    return origin_questions, questions\n\n\ndef time_replace(data):\n    \"\"\"\n    时间替换\n    :param data: 数据\n    :return: 数据\n    \"\"\"\n    question, _ = time_extract(data['question'])\n    data['question'] = question\n    if 'keywords' in data and 'values' in data['keywords']:\n        values = [time_extract(value)[0] for value in data['keywords']['values']]\n        data['keywords']['values'] = values\n    if 'sql' in data and 'conds' in data['sql']:\n        conds = [[value[0], value[1], time_extract(value[2])[0]] for value in data['sql']['conds']]\n        data['sql']['conds'] = conds\n\n\ndef stop_words():\n    nltk.download('stopwords')\n    words = stopwords.words('chinese')\n    return words\n\n\ndef cut_words_first_end(cut_word_list, content):\n    \"\"\"\n    剪切词\n    :param cut_word_list: 剪切词列表\n    :param content: 内容\n    :return: 剪切后的内容\n    \"\"\"\n    cut_result = [word for word, _ in psg.cut(content)]\n    if len(cut_result) > 1 and cut_result[0] in cut_word_list:\n        cut_result.pop(0)\n    if len(cut_result) > 1 and cut_result[-1] in cut_word_list:\n        cut_result.pop(-1)\n    return ''.join(cut_result)\n\n\ndef get_columns(table_path):\n    columns = pd.read_table(table_path, header=2)\n    return columns.columns.__array__()\n\n\ndef get_cond_op_dict():\n    cond_op_dict = {'>': 0, '<': 1, '==': 2, '!=': 3, 'like': 4, '>=': 5, '<=': 6, 'none': 7}\n    return cond_op_dict\n\n\ndef get_conn_op_dict():\n    conn_op_dict = {'none': 0, 'and': 1, 'or': 2}\n    return conn_op_dict\n\n\ndef get_agg_dict():\n    agg_dict = {'': 0, 'AVG': 1, 'MAX': 2, 'MIN': 3, 'COUNT': 4, 'SUM': 5, 'none': 6}\n    return agg_dict\n\n\ndef get_key(dict, value):\n    \"\"\"\n    根据字典的value获取key\n    :param dict: 字典\n    :param value: 值\n    :return: key\n    \"\"\"\n    return [k for k, v in dict.items() if v == value]\n\n\ndef fill_value_start_end(cond_vals, question, value, idx):\n    \"\"\"\n    :param cond_vals 待填充待值\n    :param question 问题\n    :param value 待匹配待值\n    :param 下标\n    fill [1] by the value in the question\n    结果类似[0,0,0,0,1,1,1,1,0,0,0,2,2,2,2,3,3,3,0,0,0,4,4,4,0,0,0,0,0]\n    \"\"\"\n    question_length = len(question)\n    value_length = len(value)\n    for i in range(question_length - value_length + 1):\n        if question[i:value_length + i] == value:\n            cond_vals[i: value_length + i] = [idx + 1] * value_length\n    return cond_vals\n\n\ndef count_values(cond_vals):\n    \"\"\"\n   cond_vals的值如[0,0,0,0,1,1,1,1,0,0,0,2,2,2,2,3,3,3,0,0,0,4,4,4,0,0,0,0,0]所示\n   统计出现>0的数量，连续的>0只统计一次\n   \"\"\"\n    count = 0\n    pre_value = None\n    for idx, val in enumerate(cond_vals):\n        if idx > 0:\n            pre_value = cond_vals[idx - 1]\n        if val > 0 and val != pre_value:\n            count = count + 1\n    return count\n\n\ndef get_values_name(question, cond_vals):\n    \"\"\"\n    cond_vals的值如[0,0,0,0,1,1,1,1,0,0,0,2,2,2,2,3,3,3,0,0,0,4,4,4,0,0,0,0,0]所示\n    根据cond_vals中为1的值找到question对应下标的内容\n    返回找到的内容列表，连续为1的内容作为返回列表的一个元素\n    \"\"\"\n    result = []\n    start_of_segment = None\n    for idx, current_value in enumerate(cond_vals):\n        previous_value = cond_vals[idx - 1] if idx > 0 else None\n        # 检测新的连续段的开始\n        if current_value > 0 and current_value != previous_value and start_of_segment is None:\n            start_of_segment = idx\n        # 当段结束时，添加到结果并更新段的起始位置\n        elif current_value != previous_value and start_of_segment is not None:\n            segment = question[start_of_segment:idx]\n            result.append(segment)\n            start_of_segment = None if current_value == 0 else idx\n    return result\n","metadata":{"execution":{"iopub.status.busy":"2024-01-24T03:23:38.260257Z","iopub.execute_input":"2024-01-24T03:23:38.260623Z","iopub.status.idle":"2024-01-24T03:23:38.287431Z","shell.execute_reply.started":"2024-01-24T03:23:38.260596Z","shell.execute_reply":"2024-01-24T03:23:38.286242Z"},"trusted":true},"execution_count":3,"outputs":[]},{"cell_type":"code","source":"#dataset\n# label\nclass Label(object):\n    def __init__(self, label_agg: List = None, label_conn_op=None, label_cond_cols: List = None,\n                 label_cond_ops: List = None, label_cond_vals: List = None):\n        \"\"\"\n        训练标签信息\n        :param label_agg: 聚合函数\n        :param label_conn_op: 连接操作符\n        :param label_cond_cols: 条件操列\n        :param label_cond_ops: 条件操作符\n        :param label_cond_vals: 条件值\n        \"\"\"\n        self.label_agg = label_agg\n        self.label_conn_op = label_conn_op\n        self.label_cond_ops = label_cond_ops\n        self.label_cond_cols = label_cond_cols\n        self.label_cond_vals = label_cond_vals\n\n\nclass InputFeatures(object):\n    def __init__(self, model_path=None, question_length=128, max_length=512, input_ids=None, attention_mask=None,\n                 token_type_ids=None, cls_idx=None, label: Label = None):\n        if model_path is not None:\n            self.tokenizer = BertTokenizer.from_pretrained(model_path)\n        self.question_length = question_length\n        self.max_length = max_length\n        self.input_ids = input_ids\n        self.attention_mask = attention_mask\n        self.token_type_ids = token_type_ids\n        self.cls_idx = cls_idx\n        self.label = label\n\n    def encode_expression(self, expressions: List):\n        \"\"\"\n        表达式编码\n        :param expressions: 表达式（列名或条件表达式）\n        :return: 编码后的列，及序列号（用于列与列之间的区分）\n        \"\"\"\n        encodings = self.tokenizer.batch_encode_plus(expressions)\n        expressions_encode = encodings[\"input_ids\"]\n        segment_ids = encodings[\"token_type_ids\"]\n        segment_ids = [[elem if j % 2 == 0 else 1 for elem in row] for j, row in enumerate(segment_ids)]\n        expressions_encode = [item for sublist in expressions_encode for item in sublist]\n        segment_ids = [item for sublist in segment_ids for item in sublist]\n        return torch.tensor(expressions_encode), torch.tensor(segment_ids)\n\n    def get_cls_idx(self, expressions):\n        \"\"\"\n        获取表达式标记符的位置\n        :param expressions: 表达式\n        :return:\n        \"\"\"\n        cls_idx = []\n        start = self.question_length\n        for i in range(len(expressions)):\n            cls_idx.append(int(start))\n            # 加上特殊标记的长度（例如 [CLS] 和 [SEP]）\n            start += len(expressions[i]) + 2\n        return cls_idx\n\n    def encode_question_with_expressions(self, que_length, max_length, question, expressions_encode,\n                                         expressions_segment_id):\n        \"\"\"\n        编码\n        :param que_length: 问题长度\n        :param max_length: text长度\n        :param question:  问题\n        :param expressions_encode:  编码的列\n        :param expressions_segment_id 编码的列的序列\n        :return: 编码后的text\n        \"\"\"\n\n        # 编码问题，需要填充，否则会出现长度不一致异常\n        question_encoding = self.tokenizer.encode(question, add_special_tokens=True, padding='max_length',\n                                                  max_length=que_length, truncation=True)\n\n        # 合并编码后的张量，保证张量类型(dtype)为int或long, bert的embedding的要求\n        input_ids = torch.cat([torch.tensor(question_encoding), expressions_encode], dim=0)\n        token_type_ids = torch.cat([torch.zeros(que_length, dtype=torch.long), expressions_segment_id], dim=0)\n        padding_length = max_length - len(input_ids)\n        attention_mask = torch.cat([torch.ones(len(input_ids)), torch.zeros(padding_length)], dim=0)\n        input_ids = torch.cat([input_ids, torch.zeros(padding_length, dtype=torch.long)], dim=0)\n        token_type_ids = torch.cat([token_type_ids, torch.zeros(padding_length, dtype=torch.long)], dim=0)\n\n        return input_ids, attention_mask, token_type_ids\n\n    def list_features(self, columns, datas):\n        \"\"\"\n        输入特征\n        :param columns 列\n        :param datas: 数据\n        :return: 特征信息\n        \"\"\"\n        list_features = []\n        cls_idx = self.get_cls_idx(columns)\n        expressions_encode, expressions_segment_id = self.encode_expression(columns)\n        for data in datas:\n            question = data[0]\n            # if contain label data\n            label = None\n            if len(data) > 1:\n                label = Label(label_agg=data[1], label_conn_op=data[2], label_cond_cols=data[3], label_cond_ops=data[4],\n                              label_cond_vals=data[5])\n            # 编码(question+expressions)\n            input_ids, attention_mask, token_type_ids = self.encode_question_with_expressions(self.question_length,\n                                                                                              self.max_length,\n                                                                                              question,\n                                                                                              expressions_encode,\n                                                                                              expressions_segment_id)\n            list_features.append(\n                InputFeatures(question_length=self.question_length, max_length=self.max_length, input_ids=input_ids,\n                              attention_mask=attention_mask, token_type_ids=token_type_ids, cls_idx=cls_idx,\n                              label=label))\n        return list_features\n\n\nclass Dataset(torch.utils.data.Dataset):\n    def __init__(self, features: List[InputFeatures]):\n        self.features = features\n\n    def __len__(self):\n        return len(self.features)\n\n    def __getitem__(self, item):\n        feature = self.features[item]\n        input_ids = np.array(feature.input_ids)\n        attention_mask = np.array(feature.attention_mask)\n        token_type_ids = np.array(feature.token_type_ids)\n        cls_idx = np.array(feature.cls_idx)\n        if feature.label is not None:\n            label: Label = feature.label\n            label_agg = np.array(label.label_agg)\n            label_conn_op = np.array(label.label_conn_op)\n            label_cond_cols = np.array(label.label_cond_cols)\n            label_cond_ops = np.array(label.label_cond_ops)\n            label_cond_vals = np.array([np.array(val) for val in label.label_cond_vals])\n            return input_ids, attention_mask, token_type_ids, cls_idx, label_agg, label_conn_op, label_cond_cols, label_cond_ops, label_cond_vals\n        else:\n            return input_ids, attention_mask, token_type_ids, cls_idx\n\n","metadata":{"execution":{"iopub.status.busy":"2024-01-24T03:23:39.777205Z","iopub.execute_input":"2024-01-24T03:23:39.777553Z","iopub.status.idle":"2024-01-24T03:23:39.801833Z","shell.execute_reply.started":"2024-01-24T03:23:39.777526Z","shell.execute_reply":"2024-01-24T03:23:39.800787Z"},"trusted":true},"execution_count":4,"outputs":[]},{"cell_type":"code","source":"#model\nclass ColClassifierModel(nn.Module):\n    def __init__(self, model_path, hidden_size, agg_length, conn_op_length, dropout=0.5):\n        super(ColClassifierModel, self).__init__()\n        self.bert = BertModel.from_pretrained(model_path)\n        self.dropout = nn.Dropout(dropout)\n        # out classes需要纬度必须大于label中size(classes)，否则会出现Assertion `t >= 0 && t < n_classes` failed.\n        self.agg_classifier = nn.Linear(hidden_size, agg_length)\n        self.conn_op_classifier = nn.Linear(hidden_size, conn_op_length)\n\n    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, cls_idx=None):\n        # 输出最后一层隐藏状态以及池化层\n        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)\n        dropout_output = self.dropout(outputs.pooler_output)\n        dropout_hidden_state = self.dropout(outputs.last_hidden_state)\n\n        \"\"\"\n        提取列特征信息，从dim=1即第二维中（列标记符号索引所在纬度）提取dropout_hidden_state对应该纬度的信息。\n        前提需要将cls_idx张量shape扩展成与dropout_hidden_state一致\n        \"\"\"\n        # cls_cols = dropout_hidden_state.gather(dim=1, index=cls_idx.unsqueeze(-1).expand(\n        #     dropout_hidden_state.shape[0], -1, dropout_hidden_state.shape[-1]))\n        # 简化写法\n        cls_cols = dropout_hidden_state[:, cls_idx[0], :]\n\n        out_agg = self.agg_classifier(cls_cols)\n\n        out_conn_op = self.conn_op_classifier(dropout_output)\n\n        return out_agg, out_conn_op\n\n\nclass CondClassifierModel(nn.Module):\n    def __init__(self, model_path, hidden_size, question_length, dropout=0.5):\n        super(CondClassifierModel, self).__init__()\n        self.bert = BertModel.from_pretrained(model_path)\n        self.dropout = nn.Dropout(dropout)\n        # question_length为条件最多个数\n        self.cond_cols_classifier = nn.Linear(hidden_size, question_length)\n        self.cond_ops_classifier = nn.Linear(hidden_size, question_length)\n        self.cond_vals_classifier = nn.Linear(hidden_size, question_length)\n        self.cond_count_classifier = nn.Linear(hidden_size, question_length)\n        self.question_length = question_length\n\n    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None):\n        # 输出最后一层隐藏状态\n        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)\n        dropout_output = self.dropout(outputs.pooler_output)\n        hidden_state = outputs.last_hidden_state\n\n        out_cond_count = self.cond_count_classifier(dropout_output)\n\n        # 提取问题特征信息\n        cond_values = hidden_state[:, 1:self.question_length + 1, :]\n\n        out_cond_cols = self.cond_cols_classifier(cond_values)\n        out_cond_ops = self.cond_ops_classifier(cond_values)\n        out_cond_vals = self.cond_vals_classifier(cond_values)\n\n        return out_cond_cols, out_cond_ops, out_cond_vals, out_cond_count\n","metadata":{"execution":{"iopub.status.busy":"2024-01-24T03:23:41.448195Z","iopub.execute_input":"2024-01-24T03:23:41.448569Z","iopub.status.idle":"2024-01-24T03:23:41.461586Z","shell.execute_reply.started":"2024-01-24T03:23:41.448538Z","shell.execute_reply":"2024-01-24T03:23:41.460649Z"},"trusted":true},"execution_count":5,"outputs":[]},{"cell_type":"code","source":"#train\ndef train(model: ColClassifierModel or CondClassifierModel, model_save_path, train_dataset: Dataset,\n          val_dataset: Dataset, batch_size, lr, epochs):\n    # DataLoader根据batch_size获取数据，训练时选择打乱样本\n    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n    val_loader = DataLoader(val_dataset, batch_size=batch_size)\n    # 是否使用gpu\n    use_cuda = torch.cuda.is_available()\n    device = torch.device(\"cuda:0\" if use_cuda else \"cpu\")\n    # 定义损失函数和优化器\n    criterion = nn.CrossEntropyLoss()\n    optim = Adam(model.parameters(), lr=lr)\n    if use_cuda:\n        model = model.to(device)\n        criterion = criterion.to(device)\n    best_val_avg_acc = 0\n    for epoch in range(epochs):\n        total_loss_train = 0\n        model.train()\n        # 训练进度\n        for input_ids, attention_mask, token_type_ids, cls_idx, label_agg, label_conn_op, label_cond_cols, label_cond_ops, label_cond_vals in tqdm(\n                train_loader):\n            # model要求输入的矩阵(hidden_size,sequence_size),需要把第二纬度去除.squeeze(1)\n            input_ids = input_ids.squeeze(1).to(device)\n            attention_mask = attention_mask.to(device)\n            token_type_ids = token_type_ids.squeeze(1).to(device)\n            if type(model) is ColClassifierModel:\n                # reshape(-1)合并一二纬度\n                label_agg = label_agg.to(device).reshape(-1)\n                label_conn_op = label_conn_op.to(device)\n                # 模型输出\n                out_agg, out_conn_op = model(input_ids, attention_mask, token_type_ids, cls_idx)\n                out_agg = out_agg.to(device).reshape(-1, out_agg.size(2))\n                out_conn_op = out_conn_op.to(device)\n                # 计算损失\n                loss_agg = criterion(out_agg, label_agg)\n                loss_conn_op = criterion(out_conn_op, label_conn_op)\n                # 损失比例\n                total_loss_train = loss_agg + loss_conn_op\n\n            if type(model) is CondClassifierModel:\n                label_cond_cols = label_cond_cols.to(device).reshape(-1)\n                label_cond_ops = label_cond_ops.to(device).reshape(-1)\n                label_cond_vals = label_cond_vals.to(device)\n                label_cond_count = [count_values(label_cond_val) for label_cond_val in label_cond_vals]\n                label_cond_count = torch.tensor(label_cond_count).reshape(-1).to(device)\n                label_cond_vals = label_cond_vals.reshape(-1)\n                # 模型输出\n                out_cond_cols, out_cond_ops, out_cond_vals, out_cond_count = model(input_ids, attention_mask,\n                                                                                   token_type_ids)\n                # 计算损失\n                out_cond_cols = out_cond_cols.reshape(-1, out_cond_cols.size(2))\n                out_cond_ops = out_cond_ops.reshape(-1, out_cond_ops.size(2))\n                out_cond_vals = out_cond_vals.reshape(-1, out_cond_vals.size(2))\n                lost_cond_cols = criterion(out_cond_cols, label_cond_cols)\n                lost_cond_ops = criterion(out_cond_ops, label_cond_ops)\n                lost_cond_vals = criterion(out_cond_vals, label_cond_vals)\n                lost_cond_count = criterion(out_cond_count, label_cond_count)\n                total_loss_train = ((lost_cond_cols + lost_cond_vals + lost_cond_ops) * 1 + lost_cond_count * 1)\n\n            # 模型更新\n            model.zero_grad()\n            optim.zero_grad()\n            total_loss_train.backward()\n            optim.step()\n        # 模型验证\n        val_avg_acc = 0\n        out_all_agg = []\n        out_all_conn_op = []\n        out_all_cond_cols = []\n        out_all_cond_ops = []\n        out_all_cond_vals = []\n        out_all_cond_count = []\n        label_all_agg = []\n        label_all_conn_op = []\n        label_all_cond_cols = []\n        label_all_cond_ops = []\n        label_all_cond_vals = []\n        label_all_cond_count = []\n        # 验证无需梯度计算\n        model.eval()\n        with torch.no_grad():\n            # 使用当前epoch训练好的模型验证\n            for input_ids, attention_mask, token_type_ids, cls_idx, label_agg, label_conn_op, label_cond_cols, label_cond_ops, label_cond_vals in val_loader:\n                input_ids = input_ids.squeeze(1).to(device)\n                attention_mask = attention_mask.to(device)\n                token_type_ids = token_type_ids.squeeze(1).to(device)\n                if type(model) is ColClassifierModel:\n                    label_agg = label_agg.to(device).reshape(-1)\n                    label_conn_op = label_conn_op.to(device)\n                    # 模型输出\n                    out_agg, out_conn_op = model(input_ids, attention_mask, token_type_ids, cls_idx)\n                    out_agg = out_agg.argmax(dim=2).to(device).reshape(-1)\n                    out_conn_op = out_conn_op.argmax(dim=1).to(device)\n                    out_all_agg.append(out_agg.cpu().numpy())\n                    out_all_conn_op.append(out_conn_op.cpu().numpy())\n                    label_all_agg.append(label_agg.cpu().numpy())\n                    label_all_conn_op.append(label_conn_op.cpu().numpy())\n                if type(model) is CondClassifierModel:\n                    label_cond_cols = label_cond_cols.to(device).reshape(-1)\n                    label_cond_ops = label_cond_ops.to(device).reshape(-1)\n                    label_cond_vals = label_cond_vals.to(device)\n                    label_count_value = [count_values(label_cond_val) for label_cond_val in label_cond_vals]\n                    label_cond_vals = label_cond_vals.reshape(-1)\n                    # 模型输出\n                    out_cond_cols, out_cond_ops, out_cond_vals, out_cond_count = model(input_ids, attention_mask,\n                                                                                       token_type_ids)\n                    out_cond_cols = out_cond_cols.argmax(dim=2).to(device)\n                    out_cond_ops = out_cond_ops.argmax(dim=2).to(device)\n                    out_cond_vals = out_cond_vals.argmax(dim=2).to(device)\n                    out_cond_count = out_cond_count.argmax(dim=1).to(device)\n                    out_cond_cols = out_cond_cols.reshape(-1)\n                    out_cond_ops = out_cond_ops.reshape(-1)\n                    out_cond_vals = out_cond_vals.reshape(-1)\n                    out_all_cond_cols.append(out_cond_cols.cpu().numpy())\n                    out_all_cond_ops.append(out_cond_ops.cpu().numpy())\n                    out_all_cond_vals.append(out_cond_vals.cpu().numpy())\n                    out_all_cond_count.extend(out_cond_count.cpu().numpy())\n                    label_all_cond_cols.append(label_cond_cols.cpu().numpy())\n                    label_all_cond_ops.append(label_cond_ops.cpu().numpy())\n                    label_all_cond_vals.append(label_cond_vals.cpu().numpy())\n                    label_all_cond_count.extend(label_count_value)\n\n        if type(model) is ColClassifierModel:\n            val_agg_acc = metrics.accuracy_score(np.concatenate(out_all_agg, axis=0),\n                                                 np.concatenate(label_all_agg, axis=0))\n            val_conn_op_acc = metrics.accuracy_score(np.concatenate(out_all_conn_op, axis=0),\n                                                     np.concatenate(label_all_conn_op, axis=0))\n            print(f'val_agg_acc: {val_agg_acc}')\n            print(f'val_conn_op_acc: {val_conn_op_acc}')\n            # 准确率计算逻辑\n            val_avg_acc = (val_agg_acc + val_conn_op_acc) / 2\n        if type(model) is CondClassifierModel:\n            val_cond_cols_acc = metrics.accuracy_score(np.concatenate(out_all_cond_cols, axis=0),\n                                                       np.concatenate(label_all_cond_cols, axis=0))\n            val_cond_ops_acc = metrics.accuracy_score(np.concatenate(out_all_cond_ops, axis=0),\n                                                      np.concatenate(label_all_cond_ops, axis=0))\n            val_cond_vals_acc = metrics.accuracy_score(np.concatenate(out_all_cond_vals, axis=0),\n                                                       np.concatenate(label_all_cond_vals, axis=0))\n            val_cond_count_acc = metrics.accuracy_score(label_all_cond_count, out_all_cond_count)\n            print(f'val_cond_cols_acc: {val_cond_cols_acc}')\n            print(f'val_cond_ops_acc: {val_cond_ops_acc}')\n            print(f'val_cond_vals_acc: {val_cond_vals_acc}')\n            print(f'val_cond_count_acc: {val_cond_count_acc}')\n            val_avg_acc = (val_cond_cols_acc + val_cond_ops_acc + val_cond_vals_acc + val_cond_count_acc) / 4\n            # save model\n        if val_avg_acc > best_val_avg_acc:\n            best_val_avg_acc = val_avg_acc\n            torch.save(model.state_dict(), model_save_path)\n            print(f'''best model | Val Accuracy: {best_val_avg_acc: .4f}''')\n        print(\n            f'''Epochs: {epoch + 1} \n              | Train Loss: {total_loss_train.item(): .4f} \n              | Val Accuracy: {val_avg_acc: .4f}''')\n\n\nif __name__ == '__main__':\n    hidden_size = 768\n    batch_size = 12\n    learn_rate = 2e-5\n    epochs = 100\n    question_length = 128\n    max_length = 512\n    table_path = '/kaggle/input/bert-nl2sql-train-datas/table.xlsx'\n    train_data_path = '/kaggle/input/bert-nl2sql-train-datas/train.jsonl'\n    pretrain_model_path = '/kaggle/input/bert-nl2sql-chinese-model-hgd'\n    save_column_model_path = '/kaggle/working/classifier-column-model.pkl'\n    save_value_model_path = '/kaggle/working/classifier-value-model.pkl'\n    # 读取列\n    columns = get_columns(table_path)\n    # 加载数据\n    label_datas = read_train_datas(train_data_path, question_length, columns)\n    # 提取特征数据\n    model_features = InputFeatures(pretrain_model_path, question_length, max_length).list_features(columns, label_datas)\n    # 初始化dataset\n    model_dateset = Dataset(model_features)\n    # 创建模型\n    col_model = ColClassifierModel(pretrain_model_path, hidden_size, len(get_agg_dict()), len(get_conn_op_dict()))\n    # 分割数据集\n    total_size = len(label_datas)\n    train_size = int(0.8 * total_size)\n    val_size = int(0.1 * total_size)\n    test_size = total_size - train_size - val_size\n    # 分割数据集\n    model_train_dataset, model_val_dataset, model_test_dataset = random_split(model_dateset,\n                                                                              [train_size, val_size,\n                                                                               test_size])\n    print('train column model begin')\n    train(col_model, save_column_model_path, model_train_dataset, model_val_dataset, batch_size, learn_rate,\n          epochs)\n    print('train column model finish')\n    epochs = 300\n    cond_model = CondClassifierModel(pretrain_model_path, hidden_size, question_length)\n    print('train value model begin')\n    train(cond_model, save_value_model_path, model_train_dataset, model_val_dataset, batch_size,\n          learn_rate,\n          epochs)\n    print('train value model finish')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#predict\ndef predict(columns, origin_questions, questions, predict_result_path, pretrain_model_path, column_model_path,\n            value_model_path, hidden_size, batch_size, question_length, max_length, table_name='table_name'):\n    # 创建模型\n    col_model = ColClassifierModel(pretrain_model_path, hidden_size, len(get_agg_dict()), len(get_conn_op_dict()))\n    cond_model = CondClassifierModel(pretrain_model_path, hidden_size, question_length)\n    # 提取特征数据（不含label的数据）\n    input_features = InputFeatures(pretrain_model_path, question_length, max_length).list_features(columns, questions)\n    dataset = Dataset(input_features)\n    # 预测不用打乱顺序shuffle=False\n    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n    # 是否使用gpu\n    use_cuda = torch.cuda.is_available()\n    device = torch.device(\"cuda:0\" if use_cuda else \"cpu\")\n    if use_cuda:\n        col_model = col_model.to(device)\n        cond_model = cond_model.to(device)\n    col_model.load_state_dict(torch.load(column_model_path, map_location=torch.device(device)))\n    cond_model.load_state_dict(torch.load(value_model_path, map_location=torch.device(device)))\n    # 预测\n    pre_all_agg = []\n    pre_all_conn_op = []\n    pre_all_cond_cols = []\n    pre_all_cond_ops = []\n    pre_all_cond_vals = []\n    pre_all_cond_counts = []\n    for input_ids, attention_mask, token_type_ids, cls_idx in tqdm(dataloader):\n        input_ids = input_ids.to(device)\n        attention_mask = attention_mask.to(device)\n        token_type_ids = token_type_ids.to(device)\n        out_agg, out_conn_op = col_model(input_ids, attention_mask, token_type_ids, cls_idx)\n        # 取预测结果最大值，torch.argmax找到指定纬度最大值所对应的索引（是索引，不是值）\n        pre_agg = torch.argmax(out_agg, dim=2).cpu().numpy()\n        pre_conn_op = torch.argmax(out_conn_op, dim=1).cpu().numpy()\n        pre_all_agg.extend(pre_agg)\n        pre_all_conn_op.extend(pre_conn_op)\n    for input_ids, attention_mask, token_type_ids, cls_idx in tqdm(dataloader):\n        input_ids = input_ids.to(device)\n        attention_mask = attention_mask.to(device)\n        token_type_ids = token_type_ids.to(device)\n        out_cond_cols, out_cond_ops, out_cond_vals, out_cond_count = cond_model(input_ids, attention_mask,\n                                                                                token_type_ids)\n        pre_cond_cols = torch.argmax(out_cond_cols, dim=2).cpu().numpy()\n        pre_cond_ops = torch.argmax(out_cond_ops, dim=2).cpu().numpy()\n        pre_cond_vals = torch.argmax(out_cond_vals, dim=2).cpu().numpy()\n        pre_cond_count = torch.argmax(out_cond_count, dim=1).cpu().numpy()\n        pre_all_cond_cols.extend(pre_cond_cols)\n        pre_all_cond_ops.extend(pre_cond_ops)\n        pre_all_cond_vals.extend(pre_cond_vals)\n        pre_all_cond_counts.extend(pre_cond_count)\n\n    with open(predict_result_path, 'w', encoding='utf-8') as wf:\n        for origin_question, question, agg, conn_op, cond_cols, cond_ops, cond_vals, cond_counts in zip(\n                origin_questions,\n                questions,\n                pre_all_agg,\n                pre_all_conn_op,\n                pre_all_cond_cols,\n                pre_all_cond_ops,\n                pre_all_cond_vals,\n                pre_all_cond_counts):\n            sel_col = np.where(np.array(agg) != get_agg_dict()['none'])[0]\n            agg = agg[agg != get_agg_dict()['none']]\n            cond_col = cond_cols[cond_cols <= len(columns)]\n            cond_op = cond_ops[cond_ops != get_cond_op_dict()['none']]\n            sel_col_name = [columns[idx_col] for idx_col in sel_col]\n            stop_word_list = stop_words()\n            cond_vals_name = [cut_words_first_end(stop_word_list, value_name) for value_name in\n                              get_values_name(question[0], cond_vals)]\n            conds = [[int(cond_col), int(cond_op), cond_vals_name] for\n                     _, cond_col, cond_op, cond_vals_name in zip(range(cond_counts), cond_col, cond_op, cond_vals_name)]\n            sql_dict = {\"question\": origin_question, \"table_id\": table_name,\n                        \"sql\": {\"sel\": list(map(int, sel_col)),\n                                \"agg\": list(map(int, agg)),\n                                \"limit\": 0,\n                                \"orderby\": [],\n                                \"asc_desc\": 0,\n                                \"cond_conn_op\": int(conn_op),\n                                'conds': conds},\n                        \"keywords\": {\"sel_cols\": sel_col_name, \"values\": cond_vals_name}}\n            sql_json = json.dumps(sql_dict, ensure_ascii=False)\n            wf.write(sql_json + '\\n')\n\n\nif __name__ == '__main__':\n    hidden_size = 768\n    batch_size = 12\n    question_length = 128\n    max_length = 512\n    table_path = '/kaggle/input/bert-nl2sql-train-datas/table.xlsx'\n    predict_question_path = '/kaggle/input/bert-nl2sql-train-datas/train_test.jsonl'\n    predict_result_path = '/kaggle/working/predict.jsonl'\n    pretrain_model_path = '/kaggle/input/bert-nl2sql-chinese-model-hgd'\n    column_model_path = '/kaggle/input/bert-nl2sql-result-model/classifier-column-model.pkl'\n    value_model_path = '/kaggle/input/bert-nl2sql-result-model/classifier-value-model.pkl'\n    columns = get_columns(table_path)\n    origin_questions, questions = read_predict_datas(predict_question_path)\n    predict(columns, origin_questions, questions, predict_result_path, pretrain_model_path, column_model_path,\n            value_model_path,\n            hidden_size, batch_size, question_length, max_length)","metadata":{"execution":{"iopub.status.busy":"2024-01-24T03:23:43.463443Z","iopub.execute_input":"2024-01-24T03:23:43.463779Z","iopub.status.idle":"2024-01-24T03:25:28.820474Z","shell.execute_reply.started":"2024-01-24T03:23:43.463752Z","shell.execute_reply":"2024-01-24T03:25:28.819336Z"},"trusted":true},"execution_count":6,"outputs":[{"name":"stderr","text":"Building prefix dict from the default dictionary ...\nLoading model from cache /tmp/jieba.cache\nLoading model cost 0.959 seconds.\nPrefix dict has been built successfully.\n100%|██████████| 1/1 [00:00<00:00,  2.99it/s]\n100%|██████████| 1/1 [00:00<00:00, 36.65it/s]\n","output_type":"stream"},{"name":"stdout","text":"[nltk_data] Error loading stopwords: <urlopen error [Errno -3]\n[nltk_data]     Temporary failure in name resolution>\n[nltk_data] Error loading stopwords: <urlopen error [Errno -3]\n[nltk_data]     Temporary failure in name resolution>\n[nltk_data] Error loading stopwords: <urlopen error [Errno -3]\n[nltk_data]     Temporary failure in name resolution>\n[nltk_data] Error loading stopwords: <urlopen error [Errno -3]\n[nltk_data]     Temporary failure in name resolution>\n[nltk_data] Error loading stopwords: <urlopen error [Errno -3]\n[nltk_data]     Temporary failure in name resolution>\n","output_type":"stream"}]}]}