2019BDC-keycode.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CE8B5EF4D2A443BEAAC135B7F6CE845B",
    "mdEditEnable": false
   },
   "source": [
    "### 目录\n",
    "1. 导入相关工具包\n",
    "1. 训练词向量\n",
    "1. 特征工程\n",
    "1. 模型训练"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E1E8BEBCE18D4509864206237F608F2D",
    "mdEditEnable": false
   },
   "source": [
    "## 1. 导入相关工具包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9E947ABD442F4240811CABFD239DC6FF"
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "import time\n",
    "import datetime\n",
    "import csv\n",
    "import pickle\n",
    "import random\n",
    "import os\n",
    "import numpy as np\n",
    "import lightgbm as lgb\n",
    "import pandas as pd\n",
    "import Levenshtein\n",
    "import textdistance\n",
    "import gensim\n",
    "from multiprocessing import Pool\n",
    "from fuzzywuzzy import fuzz\n",
    "from tqdm import tqdm\n",
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer\n",
    "from scipy.spatial.distance import cosine, cityblock, canberra, euclidean, minkowski, braycurtis, correlation\n",
    "from scipy.stats import skew, kurtosis\n",
    "from sklearn import linear_model\n",
    "from sklearn.externals import joblib\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from scipy.linalg import norm\n",
    "from collections import Counter, defaultdict\n",
    "random.seed(1)\n",
    "np.seterr(divide='ignore', invalid='ignore')\n",
    "tqdm.pandas(desc='My bar')\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import nn, optim\n",
    "random_seed = 2019\n",
    "torch.manual_seed(random_seed)\n",
    "\n",
    "pd.set_option('display.max_colwidth', 200)\n",
    "pd.set_option('display.max_columns', 200)\n",
    "pd.set_option('display.max_rows', 200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4D7FF138188B40839A4574FE5E127EC2",
    "mdEditEnable": false
   },
   "source": [
    "## 2. 训练词向量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BAD7A395584441BD83A1B822267DE34E"
   },
   "outputs": [],
   "source": [
    "data_path = \"/home/kesci/input/bytedance/\"\n",
    "train_path = data_path + \"train_final.csv\"\n",
    "test_path = data_path + \"test_final_part1.csv\"\n",
    "\n",
    "# base dirs\n",
    "C_OutputBase = '/home/kesci/work/output/'\n",
    "C_LocalDataBase = '/home/kesci/work/local-data/'\n",
    "C_Word2VecModelBase = '/home/kesci/work/word2vec-models/'\n",
    "\n",
    "# some name\n",
    "# test_feat_path = os.path.join(C_LocalDataBase, \"test_feat.csv\")\n",
    "test_feat_path = './data/test_feature.csv'\n",
    "C_AllText = 'text.txt'\n",
    "embedding_size = 300\n",
    "C_W2VModel = 'Word2Vec_' + str(embedding_size) + '.model'\n",
    "word2idx_path = \"./word2idx.pkl\"\n",
    "idx2word_path = \"./idx2word.pkl\"\n",
    "embedding_path = \"./pretrain_emb.npy\"\n",
    "UNK = \"<UNK>\"\n",
    "PAD = \"<PAD>\"\n",
    "# log format\n",
    "C_LogFormat = '%(asctime)s - %(levelname)s - %(message)s'\n",
    "# setting log format\n",
    "logging.basicConfig(level=logging.DEBUG, format=C_LogFormat)\n",
    "\n",
    "# 创建文件夹\n",
    "if not os.path.exists(C_OutputBase):\n",
    "    logging.debug('Create Dir ' + C_OutputBase)\n",
    "    os.makedirs(C_OutputBase)\n",
    "if not os.path.exists(C_LocalDataBase):\n",
    "    logging.debug('Create Dir ' + C_LocalDataBase)\n",
    "    os.makedirs(C_LocalDataBase)\n",
    "if not os.path.exists(C_Word2VecModelBase):\n",
    "    logging.debug('Create Dir ' + C_Word2VecModelBase)\n",
    "    os.makedirs(C_Word2VecModelBase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A0D9B3E6A84346898F8DF49838909006"
   },
   "outputs": [],
   "source": [
    "# extract title and query raw text from train and test dataset\n",
    "def extract_text(file, of, which):\n",
    "    nitem = 0\n",
    "    nitems = C_TrainLen\n",
    "    if which == 'test':\n",
    "        nitems = C_TestLen\n",
    "    \n",
    "    with open(file, 'r') as f:\n",
    "        for line in f:\n",
    "            nitem += 1\n",
    "            elements = line.strip().split(',')\n",
    "            if which == 'test':\n",
    "                assert len(elements) == 4\n",
    "            if which == 'train':\n",
    "                assert len(elements) == 5\n",
    "            of.write(elements[1])\n",
    "            of.write('\\n')\n",
    "            of.write(elements[3])\n",
    "            of.write('\\n')\n",
    "            if nitem % C_NForPrint == 0:\n",
    "                logging.info('%10d/%10d on %s' % (nitem, nitems, which))\n",
    "        \n",
    "of = open(C_LocalDataBase + C_AllText, 'w')\n",
    "\n",
    "logging.info('start to parse test dataset...')\n",
    "extract_text(C_TestFile, of, 'test')\n",
    "logging.info('start to parse train dataset...')\n",
    "extract_text(C_TrainFile, of, 'train')\n",
    "\n",
    "of.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "750B453128B045A5B7472CDAF729EE40"
   },
   "outputs": [],
   "source": [
    "# train Word2Vec model\n",
    "print(time.strftime('%F %T'))\n",
    "sentences = word2vec.LineSentence(C_LocalDataBase + C_AllText)\n",
    "w2v_model = word2vec.Word2Vec(sentences, \n",
    "    size=C_W2VSize, window=C_W2VWindow, min_count=C_W2VMinCount, workers=C_W2VWorkers, sg=1)\n",
    "w2v_model.save(C_Word2VecModelBase + C_W2VModel_300)\n",
    "print(time.strftime('%F %T'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B8854ADED2244AEC83A63F3F72EB4C6C",
    "mdEditEnable": false
   },
   "source": [
    "### 测试词向量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "B02D9B68271B44D09CD17C379753647A"
   },
   "outputs": [],
   "source": [
    "# test Word2Vec model\n",
    "w2v_model = gensim.models.Word2Vec.load(C_Word2VecModelBase + C_W2VModel_300)\n",
    "print(w2v_model['1'])\n",
    "print(w2v_model.most_similar(['1']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DA3F3C60313D44FC95EB2D3010269DFB",
    "mdEditEnable": false
   },
   "source": [
    "## 3. 特征工程"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0834758195644EB39C937342C93E28E3",
    "mdEditEnable": false
   },
   "source": [
    "### 特征提取及辅助函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5D55A40535914C75B5730E11706FA3F0"
   },
   "outputs": [],
   "source": [
    "# 获取公共词的数量\n",
    "def q_t_common_words(query, title):\n",
    "    query = set(query)\n",
    "    title = set(title)\n",
    "    return len(query & title)\n",
    "\n",
    "# 计算离散的词的jaccard距离\n",
    "def jaccard(query, title):\n",
    "    query = set(query)\n",
    "    title = set(title)\n",
    "    q_t_intersection_len = len(query & title)\n",
    "    q_t_union_len = len(query | title)\n",
    "    return q_t_intersection_len / q_t_union_len\n",
    "\n",
    "# 计算离散的词的相似度\n",
    "def query_title_similarity(query, title, name):\n",
    "    outlier_sample = 0\n",
    "    try:\n",
    "        counter = CountVectorizer(analyzer='word', token_pattern=u\"(?u)\\\\b\\\\w+\\\\b\")\n",
    "        counter.fit([query, title])\n",
    "        result = counter.transform([query, title]).toarray()\n",
    "        vec1, vec2 = result[0], result[1]\n",
    "        if name == 'euclid':\n",
    "            return np.linalg.norm(vec1 - vec2)  # euclid_dis\n",
    "        elif name == 'manhattan':\n",
    "            return np.sum(np.abs(vec1 - vec2))  # manhattan_dis\n",
    "        else:\n",
    "            return np.sum(vec1 * vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))  # cosine\n",
    "    except Exception as e:\n",
    "        print('Meet a outlier sample!')\n",
    "        outlier_sample += 1\n",
    "        return 0\n",
    "        \n",
    "# 加载word2vec模型\n",
    "w2v_model = gensim.models.Word2Vec.load('/home/kesci/work/word2vec-models/Word2Vec_300.model')\n",
    "\n",
    "# 计算词向量的相似度\n",
    "def get_w2v(query, title, num):\n",
    "    q = np.zeros(300)\n",
    "    count = 0\n",
    "    for w in query:\n",
    "        if w in w2v_model.wv:\n",
    "            q += w2v_model.wv[w]\n",
    "            count += 1\n",
    "    if count == 0:\n",
    "        query_vec = q\n",
    "    query_vec = (q/count).tolist()\n",
    "    \n",
    "    t = np.zeros(300)\n",
    "    count = 0\n",
    "    for w in title:\n",
    "        if w in w2v_model.wv:\n",
    "            t += w2v_model.wv[w]\n",
    "            count += 1\n",
    "    if count == 0:\n",
    "        title_vec = q\n",
    "    title_vec = (t/count).tolist()\n",
    "    \n",
    "    if num == 1:\n",
    "        try:\n",
    "            vec_cosine = cosine(query_vec, title_vec)\n",
    "            return vec_cosine\n",
    "        except Exception as e:\n",
    "            return 0\n",
    "    if num == 2:\n",
    "        try:\n",
    "            vec_canberra = canberra(query_vec, title_vec) / len(query_vec)\n",
    "            return vec_canberra\n",
    "        except Exception as e:\n",
    "            return 0\n",
    "    if num == 3:\n",
    "        try:\n",
    "            vec_cityblock = cityblock(query_vec, title_vec) / len(query_vec)\n",
    "            return vec_cityblock\n",
    "        except Exception as e:\n",
    "            return 0\n",
    "    if num == 4:\n",
    "        try:\n",
    "            vec_euclidean = euclidean(query_vec, title_vec)\n",
    "            return vec_euclidean\n",
    "        except Exception as e:\n",
    "            return 0\n",
    "    if num == 5:\n",
    "        try:\n",
    "            vec_braycurtis = braycurtis(query_vec, title_vec)\n",
    "            return vec_braycurtis\n",
    "        except Exception as e:\n",
    "            return 0\n",
    "    if num == 6:\n",
    "        try:\n",
    "            vec_minkowski = minkowski(query_vec, title_vec)\n",
    "            return vec_minkowski\n",
    "        except Exception as e:\n",
    "            return 0\n",
    "    if num == 7:\n",
    "        try:\n",
    "            vec_correlation = correlation(query_vec, title_vec)\n",
    "            return vec_correlation\n",
    "        except Exception as e:\n",
    "            return 0\n",
    "\n",
    "# 计算n-gram特征  \n",
    "def ngram_nums(query, title, n=2):\n",
    "    try:\n",
    "        count_vec = CountVectorizer(ngram_range=(n, n), analyzer='word', token_pattern=u\"(?u)\\\\b\\\\w+\\\\b\")\n",
    "        analyzer = count_vec.build_analyzer()\n",
    "        return 1 if len(set(analyzer(query)) & set(analyzer(title))) > 0 else 0\n",
    "    except Exception as e:\n",
    "        return 0\n",
    "\n",
    "# 提取特征的函数\n",
    "def extract_feature(data):\n",
    "\n",
    "    time1 = datetime.datetime.now()\n",
    "    print('Extracting features start.')\n",
    "    \n",
    "    data['query_split_30'] = data['query'].apply(lambda x: x.split())\n",
    "    data['title_split_30'] = data['title'].apply(lambda x: x.split())\n",
    "    \n",
    "    # 词向量的相似度特征\n",
    "    data['vec_cosine'] = data.progress_apply(lambda index: get_w2v(index['query_split_30'], index['title_split_30'], 1), axis=1)\n",
    "    data['vec_canberra'] = data.progress_apply(lambda index: get_w2v(index['query_split_30'], index['title_split_30'], 2), axis=1)\n",
    "    data['vec_cityblock'] = data.progress_apply(lambda index: get_w2v(index['query_split_30'], index['title_split_30'], 3), axis=1)\n",
    "    data['vec_euclidean'] = data.progress_apply(lambda index: get_w2v(index['query_split_30'], index['title_split_30'], 4), axis=1)\n",
    "    data['vec_braycurtis'] = data.progress_apply(lambda index: get_w2v(index['query_split_30'], index['title_split_30'], 5), axis=1)\n",
    "    # data['vec_minkowski'] = data.progress_apply(lambda index: get_w2v(index['query_split_30'], index['title_split_30'], 6), axis=1)\n",
    "    data['vec_correlation'] = data.progress_apply(lambda index: get_w2v(index['query_split_30'], index['title_split_30'], 7), axis=1)\n",
    "    \n",
    "    data['vec_cosine'] = data['vec_cosine'].astype('float32')\n",
    "    data['vec_canberra'] = data['vec_canberra'].astype('float32')\n",
    "    data['vec_cityblock'] = data['vec_cityblock'].astype('float32')\n",
    "    data['vec_euclidean'] = data['vec_euclidean'].astype('float32')\n",
    "    data['vec_braycurtis'] = data['vec_braycurtis'].astype('float32')\n",
    "    data['vec_correlation'] = data['vec_correlation'].astype('float32')\n",
    "    \n",
    "    data.drop(['query_split_30', 'title_split_30'], axis=1, inplace=True)\n",
    "    gc.collect()\n",
    "    \n",
    "    # 统计特征\n",
    "    data['query_split'] = data['query'].apply(lambda x: x.split())\n",
    "    data['title_split'] = data['title'].apply(lambda x: x.split())\n",
    "    data['query_len'] = data['query_split'].map(len)\n",
    "    data['title_len'] = data['title_split'].map(len)\n",
    "    data['q_t_rate'] = data['query_len'] / data['title_len']\n",
    "    data['q_t_common_words'] = data.apply(lambda index: q_t_common_words(index.query_split, index.title_split), axis=1)\n",
    "    # data['qlen_gt_tlen'] = data.apply(lambda index: 1 if index.query_len > index.title_len else 0, axis=1)\n",
    "    data['common_words_ql_rate'] = data['q_t_common_words'] / data['query_len']\n",
    "    data['common_words_tl_rate'] = data['q_t_common_words'] / data['title_len']\n",
    "    \n",
    "    data.drop(['query_split', 'title_split'], axis=1, inplace=True)\n",
    "    gc.collect()\n",
    "    \n",
    "    data['title_len'] = data['title_len'].astype('int16')\n",
    "    data['q_t_rate'] = data['q_t_rate'].astype('float32')\n",
    "    data['q_t_common_words'] = data['q_t_common_words'].astype('int16')\n",
    "    data['common_words_ql_rate'] = data['common_words_ql_rate'].astype('float32')\n",
    "    \n",
    "    # 多项式特征\n",
    "    poly = PolynomialFeatures(2) # 设置多项式阶数为2，其它默认\n",
    "    ploy_features = poly.fit_transform(data[['common_words_ql_rate', 'common_words_tl_rate']].values)\n",
    "    ploy_features = pd.DataFrame(ploy_features[:, 3:4], columns=['q_in_t_n_ql_r_sq'])#, 'q_in_t_n_tl_r_sq', 'q_in_t_num_ql*tl_rate'\n",
    "    data = pd.concat([data, ploy_features], axis=1)\n",
    "    data['q_in_t_n_ql_r_sq'] = data['q_in_t_n_ql_r_sq'].astype('float32')\n",
    "    \n",
    "    data.drop(['common_words_tl_rate', 'query_len'], axis=1, inplace=True)\n",
    "    \n",
    "    # 统计特征\n",
    "    query_cnt = data.groupby(['query_id'], as_index=False)['query'].count().rename(columns={'query': 'query_cnt'})\n",
    "    data = pd.merge(data, query_cnt, on='query_id', how='left')\n",
    "    title_cnt = data['title'].value_counts().reset_index(name='title_cnt').rename(columns={'index': 'title'})\n",
    "    data = pd.merge(data, title_cnt, on='title', how='left')\n",
    "    del query_cnt#, title_cnt\n",
    "    gc.collect()\n",
    "\n",
    "    title_len_min = data.groupby(['query_id'], as_index=False)['title_len'].agg(min).rename(columns={'title_len': 'title_len_min'})\n",
    "    data = pd.merge(data, title_len_min, on='query_id', how='left')\n",
    "    del title_len_min #title_len_max, title_len_mean\n",
    "    gc.collect()\n",
    "    \n",
    "    data['query_cnt'] = data['query_cnt'].astype('int16')\n",
    "    data['title_len_min'] = data['title_len_min'].astype('int16')\n",
    "    \n",
    "    # 相似性特征\n",
    "    data['levenshtein_distance'] = data.apply(lambda line: Levenshtein.distance(line['query'], line['title']), axis=1)\n",
    "    data['levenshtein_ratio'] = data.apply(lambda line: Levenshtein.ratio(line['query'], line['title']), axis=1)\n",
    "    data['levenshtein_jaro'] = data.apply(lambda line: Levenshtein.jaro(line['query'], line['title']), axis=1)\n",
    "    data['levenshtein_jaro_winkler'] = data.apply(lambda line: Levenshtein.jaro_winkler(line['query'], line['title']), axis=1)\n",
    "    \n",
    "    data['levenshtein_ratio'] = data['levenshtein_ratio'].astype('float32')\n",
    "    data['levenshtein_jaro'] = data['levenshtein_jaro'].astype('float32')\n",
    "    data['levenshtein_jaro_winkler'] = data['levenshtein_jaro_winkler'].astype('float32')\n",
    "    \n",
    "    data['leven_dis_rank'] = data.groupby(['query'], as_index=False)['levenshtein_distance'].rank(method='average', ascending=True)\n",
    "    data['hamming_similarity'] = data.apply(lambda line: textdistance.Hamming(qval=None).similarity(line['query'], line['title']), axis=1)\n",
    "    data['hamming_normalized_distance'] = data.apply(lambda line: textdistance.Hamming(qval=None).normalized_distance(line['query'], line['title']), axis=1)\n",
    "\n",
    "    data['leven_dis_rank'] = data['leven_dis_rank'].astype('float32')\n",
    "    data['hamming_similarity'] = data['hamming_similarity'].astype('int32')\n",
    "    data['hamming_normalized_distance'] = data['hamming_normalized_distance'].astype('float32')\n",
    "    \n",
    "    data['levenshtein_distance'] = data.apply(lambda line: textdistance.Levenshtein(qval=None).distance(line['query'], line['title']), axis=1)\n",
    "    data['levenshtein_similarity'] = data.apply(lambda line: textdistance.Levenshtein(qval=None).similarity(line['query'], line['title']), axis=1)\n",
    "    data['levenshtein_similarity_rank'] = data.groupby(['query'], as_index=False)['levenshtein_similarity'].rank(method='average', ascending=True)\n",
    "\n",
    "    data['jaccard_similarity'] = data.apply(lambda line: textdistance.Jaccard(qval=None).similarity(line['query'], line['title']), axis=1)\n",
    "    data['jaccard_similarity_rank'] = data.groupby(['query'], as_index=False)['jaccard_similarity'].rank(method='average', ascending=True)\n",
    "    data.drop(['jaccard_similarity'], axis=1, inplace=True)\n",
    "    gc.collect()\n",
    "    \n",
    "    data['levenshtein_distance'] = data['levenshtein_distance'].astype('int32')\n",
    "    data['levenshtein_similarity'] = data['levenshtein_similarity'].astype('int32')\n",
    "    data['levenshtein_similarity_rank'] = data['levenshtein_similarity_rank'].astype('float32')\n",
    "    data['jaccard_similarity_rank'] = data['jaccard_similarity_rank'].astype('float32')\n",
    "\n",
    "    data.drop(['query', 'title'], axis=1, inplace=True)#'jaccard_dis' \n",
    "    gc.collect()\n",
    "    \n",
    "    time2 = datetime.datetime.now()\n",
    "    print('Extracting features end.')\n",
    "    print('Cost time: {}s'.format((time2 - time1).seconds))\n",
    "    print('After feature engineering dataset size:\\n{}'.format(data.shape))\n",
    "    return data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FF8BA81ACF0A409A8901DE69E6305986",
    "mdEditEnable": false
   },
   "source": [
    "### 提取特征并保存特征文件到磁盘"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CEBD3B325A354A1B8C42ACF91B61BE0F"
   },
   "outputs": [],
   "source": [
    "# 设置数据类型\n",
    "orin_types = {'query_id': np.int32, 'query_title_id': np.int16, 'label': np.int8}\n",
    "# 数据加载函数（这里加载的是第6亿9600万到第7亿的400万数据）\n",
    "def load_train_data(rows=2000000):\n",
    "    train = pd.read_csv('/home/kesci/input/bytedance/train_final.csv', encoding='utf-8', \n",
    "                        names=['query_id', 'query', 'query_title_id', 'title', 'label'], skiprows=696000000, \n",
    "                        nrows=4000000, dtype=orin_types)\n",
    "    print('Training set size:{}'.format(train.shape))\n",
    "    return train\n",
    "    \n",
    "# 开始特征提取\n",
    "time_now = datetime.datetime.now().strftime('%Y/%m/%d-%H:%M:%S')\n",
    "print('Time: {}'.format(time_now))\n",
    "start = time.clock()\n",
    "print('Data loading ...')\n",
    "train = load_train_data()   # 5000000  10000000, validate\n",
    "print('=' * 50)\n",
    "\n",
    "# 保存特征文件（这里展示的是第6亿9600万到第7亿的400万数据）\n",
    "train_feature_696_70000w = extract_feature(train.copy())\n",
    "train_feature_696_70000w.to_csv('/home/kesci/work/data/ml/train_feature_696_70000w.csv', index=None)\n",
    "del train_feature_696_70000w, train\n",
    "gc.collect()\n",
    "\n",
    "print('Train features extracted over.')\n",
    "print('-' * 50)\n",
    "time_now = datetime.datetime.now().strftime('%Y/%m/%d-%H:%M:%S')\n",
    "print('Time: {}'.format(time_now))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "300101A4E9EC4213AAAC3D247C569FB9",
    "mdEditEnable": false
   },
   "source": [
    "## 4. 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ECFD47F2F4D745FE8CA6F8862838FE60"
   },
   "outputs": [],
   "source": [
    "# 加载word2vec模型\n",
    "logging.basicConfig(level=logging.DEBUG, format=C_LogFormat)\n",
    "w2v_model = gensim.models.Word2Vec.load(C_Word2VecModelBase + C_W2VModel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DB46EE451AC040B48C6CBEA96F135E9B"
   },
   "outputs": [],
   "source": [
    "# 构建词表并保存\n",
    "def build_vocab(train_path, min_count=5):\n",
    "    count = [(UNK, -1), (PAD, -1)]\n",
    "    counter = defaultdict(int)\n",
    "    csvfile = open(train_path, newline='')\n",
    "    spamreader = csv.reader(csvfile)\n",
    "    for i, row in enumerate(spamreader): \n",
    "        query_id, query, query_title_id, title, label = row\n",
    "        for w in query.split():\n",
    "            counter[w] += 1\n",
    "        for w in title.split():\n",
    "            counter[w] += 1\n",
    "        if i % 10000000 == 0:\n",
    "            print(i, len(counter))\n",
    "        if i > 10000 * 10000:\n",
    "            break\n",
    "            \n",
    "    for w in counter:\n",
    "        c = counter[w]\n",
    "        if c >= min_count:\n",
    "            count.append((w, c))\n",
    "    word2idx = dict()\n",
    "    for word, _ in count:\n",
    "        word2idx[word] = len(word2idx)\n",
    "    idx2word = dict(zip(word2idx.values(), word2idx.keys()))\n",
    "    return word2idx, idx2word\n",
    "\n",
    "word2idx, idx2word = build_vocab(train_path, min_count=7)\n",
    "pickle.dump(word2idx, open(word2idx_path,'wb'))\n",
    "pickle.dump(idx2word, open(idx2word_path, 'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "097C07B3D1F44B669EA2E5669D4B6AA2"
   },
   "outputs": [],
   "source": [
    "# 根据词表和预训练word2vec构建word embedding矩阵并保存\n",
    "pretrain_emb = []\n",
    "vocab_size = len(idx2word)\n",
    "print(\"index2word len:\", vocab_size)\n",
    "embedding_size = w2v_model.layer1_size\n",
    "print(\"embedding size:\", embedding_size)\n",
    "count = 0\n",
    "for i in range(vocab_size):\n",
    "    word = idx2word[i]\n",
    "    if word in w2v_model.wv:\n",
    "        pretrain_emb.append(w2v_model.wv[word])\n",
    "    else:\n",
    "        count += 1\n",
    "        if word == PAD:\n",
    "            rand_emb = np.zeros(embedding_size)\n",
    "        else:\n",
    "            rand_emb = np.random.normal(loc=0, scale=1, size=embedding_size)\n",
    "        pretrain_emb.append(rand_emb)\n",
    "print(\"rand init count:\",count)\n",
    "pretrain_emb = np.array(pretrain_emb)\n",
    "np.save(embedding_path, pretrain_emb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "45AE29040FE147E48A95BA02F093649E"
   },
   "outputs": [],
   "source": [
    "# 指定验证集和全体数据集的大小以及相应的特征文件\n",
    "valid_size = 100 * 10000\n",
    "all_data_size = 10000 * 10000\n",
    "train_size = all_data_size - valid_size\n",
    "feat_path = \"./data/train_feature_1E.csv\"\n",
    "# 训练第一亿skip_size为0，训练第二亿skip_size为1亿，以此类推\n",
    "skip_size = 0 * 10000 * 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F1F421840E8249599056A9D6A0444621"
   },
   "outputs": [],
   "source": [
    "output_size = 2\n",
    "feat_dim = 24 # 37\n",
    "        \n",
    "class ESIM_LSTM(nn.Module):\n",
    "    def __init__(self, pretrain=True, is_training=True):\n",
    "        super(ESIM_LSTM, self).__init__()\n",
    "        self.vocab_size = vocab_size\n",
    "        self.hidden_dim = 128\n",
    "        self.feat_dim = feat_dim\n",
    "        self.feat_hidden_dim = 100\n",
    "        self.merge_dim = 300\n",
    "        self.dropout_rate = 0.2\n",
    "        self.embedding_dim = embedding_size\n",
    "        if pretrain:\n",
    "            self.word_embeds = nn.Embedding(vocab_size, self.embedding_dim)\n",
    "            if os.path.exists(embedding_path) and is_training:\n",
    "                print(\"Loading pretrain embedding...\")\n",
    "                self.word_embeds.weight.data.copy_(torch.from_numpy(np.load(embedding_path))) \n",
    "                self.word_embeds.weight.requires_grad = False\n",
    "                print(\"Loaded\")\n",
    "        else:\n",
    "            self.word_embeds = nn.Embedding(vocab_size, self.embedding_dim)\n",
    "        self.q_encoder = nn.LSTM(self.embedding_dim, hidden_size=self.hidden_dim,\n",
    "                          num_layers=1,\n",
    "                          dropout=self.dropout_rate,\n",
    "                          bidirectional=True, batch_first=True)\n",
    "        \n",
    "        self.composition_layer = nn.LSTM(input_size=4*2*self.hidden_dim,\n",
    "                            hidden_size=self.hidden_dim,\n",
    "                            num_layers=1,\n",
    "                            batch_first=True,\n",
    "                            bidirectional=True)\n",
    "        self.esim_fc = nn.Linear(8*self.hidden_dim, 4*self.hidden_dim)\n",
    "        self.feat_fc = nn.Linear(self.feat_dim, self.feat_hidden_dim)\n",
    "        self.merge_fc = nn.Linear(4*self.hidden_dim+self.feat_hidden_dim, self.merge_dim)\n",
    "        self.final_fc = nn.Linear(self.merge_dim, output_size)\n",
    "        self.dropout = nn.Dropout(p=self.dropout_rate)\n",
    "        \n",
    "    def forward(self, query, title, feat):\n",
    "        def _attention(q1, q2):\n",
    "            # align_mat: [batch, q1_len, q2_len]\n",
    "            align_mat = torch.bmm(q1, q2.transpose(1, 2))\n",
    "            # [batch, q1_len, q2_feature_dim]\n",
    "            align_a = torch.bmm(F.softmax(align_mat, dim=2), q2)\n",
    "            # [batch, q2_len, q1_feature_dim]\n",
    "            align_b = torch.bmm(F.softmax(align_mat, dim=1).transpose(1, 2), q1)\n",
    "            return align_a, align_b\n",
    "\n",
    "        def _pooling(inputs):\n",
    "            max_output, _ = torch.max(inputs, 1)\n",
    "            avg_output = torch.mean(inputs, 1)\n",
    "            return torch.cat([max_output,avg_output], 1)\n",
    "        # In: batch_size * q_len\n",
    "        # Out: batch_size * q_len * word_embedding_dim\n",
    "        q1_embed = self.word_embeds(query)\n",
    "        q2_embed = self.word_embeds(title)\n",
    "        # q1_encoding: [batch_size,q1_len,output_size]\n",
    "        # q2_encoding: [batch_size,q2_len,output_size]\n",
    "        q1_encoding, _ = self.q_encoder(q1_embed)\n",
    "        q2_encoding, _ = self.q_encoder(q2_embed)\n",
    "        q1_encoding = self.dropout(q1_encoding)\n",
    "        q2_encoding = self.dropout(q2_encoding)\n",
    "        # align_a: [batch, q1_len, q2_feature_dim]\n",
    "        # align_b: [batch, q2_len, q1_feature_dim]\n",
    "        align_a, align_b = _attention(q1_encoding, q2_encoding)\n",
    "        # m_a: [batch, q1_len, 3*q1_feature_dim]\n",
    "        # m_b: [batch, q2_len, 3*q2_feature_dim]\n",
    "        m_a = torch.cat([q1_encoding, align_a, torch.abs(q1_encoding-align_a), q1_encoding*align_a], 2)\n",
    "        m_b = torch.cat([q2_encoding, align_b, torch.abs(q2_encoding-align_b), q2_encoding*align_b], 2)\n",
    "        v_a, _ = self.composition_layer(m_a)\n",
    "        v_b, _ = self.composition_layer(m_b)\n",
    "        \n",
    "        v_a = _pooling(v_a)\n",
    "        v_b = _pooling(v_b)\n",
    "        v = torch.cat([v_a, v_b], 1)\n",
    "        v = self.dropout(v)\n",
    "        esim_output = F.leaky_relu(self.esim_fc(v))\n",
    "        feat_output = F.leaky_relu(self.feat_fc(feat))\n",
    "        merge_output = torch.cat([esim_output, feat_output], dim=1)\n",
    "        output = self.final_fc(F.leaky_relu(self.merge_fc(merge_output)))\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5A6DE47F6D3846D1B322D2681B17C53B"
   },
   "outputs": [],
   "source": [
    "# 设置模型训练相关参数，指定待训练的模型\n",
    "batch_size = 512\n",
    "epoch_num = 15\n",
    "use_cuda = True\n",
    "\n",
    "begin_epoch = 0\n",
    "model = ESIM_LSTM(pretrain=True, is_training=True)\n",
    "\n",
    "if use_cuda:\n",
    "    model = model.cuda()\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "init_lr = 0.0002\n",
    "weight_decay = 0.0\n",
    "print_step = 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "75F0E18A015040F08947B6857BD9474D"
   },
   "outputs": [],
   "source": [
    "# 模型训练与验证，并保存验证得分最高的模型\n",
    "optimizer = optim.Adam(model.parameters(), lr=init_lr, weight_decay=weight_decay)\n",
    "max_qauc = 0\n",
    "for epoch in range(begin_epoch, begin_epoch+epoch_num):\n",
    "    running_loss = 0.0\n",
    "    count = 0\n",
    "    model = model.train()\n",
    "    nowTime=datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n",
    "    print(\"Begin | Time: %s, Epoch: %d\" % (nowTime, epoch))\n",
    "    for query, title, feat, y in gen_batch_data(train_path, batch_size):\n",
    "        optimizer.zero_grad()\n",
    "        if use_cuda:\n",
    "            query, title, feat, y = query.cuda(), title.cuda(), feat.cuda(), y.cuda()\n",
    "        output = model(query, title, feat)\n",
    "        loss = loss_fn(output, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        cur_loss = loss.item()\n",
    "        running_loss += cur_loss\n",
    "        count += 1\n",
    "        if count % print_step == 0:\n",
    "            nowTime=datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n",
    "            print(\"Time: %s, Epoch: %d, Count: %d, Loss: %.4f\" % (nowTime, epoch, count, running_loss/count))\n",
    "    lr = max(0.0001, (init_lr * (0.9)**(epoch+1-begin_epoch)))\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "    nowTime=datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n",
    "    print(\"END | Time: %s, Epoch: %d, Loss: %.4f, lr: %.5f\" % (nowTime, epoch, running_loss/count, lr))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}