{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp datasets.criteo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Criteo\n",
    "> Criteo dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *\n",
    "from fastcore.nb_imports import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import math\n",
    "import shutil\n",
    "import struct\n",
    "from collections import defaultdict\n",
    "from functools import lru_cache\n",
    "from pathlib import Path\n",
    "\n",
    "import lmdb\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch.utils.data\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import Dataset, TensorDataset\n",
    "\n",
    "from recohut.datasets.bases.common import Dataset\n",
    "from recohut.utils.common_utils import download_url"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Criteo sample dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class CriteoSampleDataset(Dataset):\n",
    "    \"\"\"Criteo Sample Dataset\n",
    "\n",
    "    Reference:\n",
    "        1. https://github.com/huangjunheng/recommendation_model/blob/master/DCN/dcn.py\n",
    "    \"\"\"\n",
    "    url = 'https://github.com/RecoHut-Datasets/criteo/raw/v1/dac_sample.txt'\n",
    "\n",
    "    def __init__(self, root, test_size=0.2, random_seed=42):\n",
    "        super().__init__(root)\n",
    "        self.test_size = test_size\n",
    "        self.random_seed = random_seed\n",
    "\n",
    "        self._process()\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        return 'dac_sample.txt'\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        return ['train.pt', 'test.pt']\n",
    "\n",
    "    def download(self):\n",
    "        path = download_url(self.url, self.raw_dir)\n",
    "\n",
    "    def process(self):\n",
    "        sparse_feature = ['C' + str(i) for i in range(1, 27)]\n",
    "        dense_feature = ['I' + str(i) for i in range(1, 14)]\n",
    "        col_names = ['label'] + dense_feature + sparse_feature\n",
    "        data = pd.read_csv(self.raw_paths[0], names=col_names, sep='\\t')\n",
    "\n",
    "        data[sparse_feature] = data[sparse_feature].fillna('-1', )\n",
    "        data[dense_feature] = data[dense_feature].fillna('0',)\n",
    "\n",
    "        feat_sizes = {}\n",
    "        feat_sizes_dense = {feat:1 for feat in dense_feature}\n",
    "        feat_sizes_sparse = {feat:len(data[feat].unique()) for feat in sparse_feature}\n",
    "        feat_sizes.update(feat_sizes_dense)\n",
    "        feat_sizes.update(feat_sizes_sparse)\n",
    "        self.feat_sizes = feat_sizes\n",
    "\n",
    "        from sklearn.preprocessing import LabelEncoder, MinMaxScaler\n",
    "        for feat in sparse_feature:\n",
    "            lbe = LabelEncoder()\n",
    "            data[feat] = lbe.fit_transform(data[feat])\n",
    "        nms = MinMaxScaler(feature_range=(0, 1))\n",
    "        data[dense_feature] = nms.fit_transform(data[dense_feature])\n",
    "\n",
    "        fixlen_feature_columns = [(feat,'sparse') for feat in sparse_feature ]  + [(feat,'dense') for feat in dense_feature]\n",
    "        self.dnn_feature_columns = fixlen_feature_columns\n",
    "        self.linear_feature_columns = fixlen_feature_columns\n",
    "\n",
    "        from sklearn.model_selection import train_test_split\n",
    "        train, test = train_test_split(data, test_size=self.test_size, random_state=self.random_seed)\n",
    "\n",
    "        train_label = pd.DataFrame(train['label'])\n",
    "        train = train.drop(columns=['label'])\n",
    "        train_tensor_data = TensorDataset(torch.from_numpy(np.array(train)), torch.from_numpy(np.array(train_label)))\n",
    "        torch.save(train_tensor_data, self.processed_paths[0])\n",
    "\n",
    "        test_label = pd.DataFrame(test['label'])\n",
    "        test = test.drop(columns=['label'])\n",
    "        test_tensor_data = TensorDataset(torch.from_numpy(np.array(test)), torch.from_numpy(np.array(test_label)))\n",
    "        torch.save(test_tensor_data, self.processed_paths[1])\n",
    "\n",
    "    def load(self):\n",
    "        train_tensor_data = torch.load(self.processed_paths[0])\n",
    "        test_tensor_data = torch.load(self.processed_paths[1])\n",
    "        return train_tensor_data, test_tensor_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n",
      "Done!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch/epoches: 0/20, train loss: 0.536, test auc: 0.679\n",
      "epoch/epoches: 1/20, train loss: 0.499, test auc: 0.702\n",
      "epoch/epoches: 2/20, train loss: 0.486, test auc: 0.722\n",
      "epoch/epoches: 3/20, train loss: 0.477, test auc: 0.740\n",
      "epoch/epoches: 4/20, train loss: 0.468, test auc: 0.745\n",
      "epoch/epoches: 5/20, train loss: 0.463, test auc: 0.752\n",
      "epoch/epoches: 6/20, train loss: 0.460, test auc: 0.757\n",
      "epoch/epoches: 7/20, train loss: 0.458, test auc: 0.753\n",
      "epoch/epoches: 8/20, train loss: 0.455, test auc: 0.758\n",
      "epoch/epoches: 9/20, train loss: 0.452, test auc: 0.759\n",
      "epoch/epoches: 10/20, train loss: 0.450, test auc: 0.757\n",
      "epoch/epoches: 11/20, train loss: 0.449, test auc: 0.758\n",
      "epoch/epoches: 12/20, train loss: 0.444, test auc: 0.758\n",
      "epoch/epoches: 13/20, train loss: 0.438, test auc: 0.757\n",
      "epoch/epoches: 14/20, train loss: 0.428, test auc: 0.753\n",
      "epoch/epoches: 15/20, train loss: 0.413, test auc: 0.750\n",
      "epoch/epoches: 16/20, train loss: 0.395, test auc: 0.744\n",
      "epoch/epoches: 17/20, train loss: 0.381, test auc: 0.735\n",
      "epoch/epoches: 18/20, train loss: 0.365, test auc: 0.737\n",
      "epoch/epoches: 19/20, train loss: 0.354, test auc: 0.735\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "from sklearn.metrics import log_loss, roc_auc_score\n",
    "\n",
    "from recohut.models.dcn import DCNv2 as DCN\n",
    "\n",
    "\n",
    "def get_auc(loader, model):\n",
    "    pred, target = [], []\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for x, y in loader:\n",
    "            x, y = x.to(device).float(), y.to(device).float()\n",
    "            y_hat = model(x)\n",
    "            pred += list(y_hat.cpu().numpy())\n",
    "            target += list(y.cpu().numpy())\n",
    "    auc = roc_auc_score(target, pred)\n",
    "    return auc\n",
    "\n",
    "\n",
    "root = '/content/data'\n",
    "batch_size = 1024\n",
    "lr = 1e-2\n",
    "wd = 1e-3\n",
    "epoches = 20\n",
    "seed = 2022\n",
    "embedding_size = 4\n",
    "device = 'cpu'\n",
    "\n",
    "ds = CriteoSampleDataset(root=root)\n",
    "train_tensor_data, test_tensor_data = ds.load()\n",
    "train_loader = DataLoader(train_tensor_data, shuffle=True, batch_size=batch_size)\n",
    "test_loader = DataLoader(test_tensor_data, batch_size=batch_size)\n",
    "\n",
    "model = DCN(ds.feat_sizes, embedding_size, ds.linear_feature_columns, ds.dnn_feature_columns).to(device)\n",
    "loss_func = nn.BCELoss(reduction='mean')\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=wd)\n",
    "\n",
    "for epoch in range(epoches):\n",
    "    total_loss_epoch = 0.0\n",
    "    total_tmp = 0\n",
    "    model.train()\n",
    "    for index, (x, y) in enumerate(train_loader):\n",
    "        x, y = x.to(device).float(), y.to(device).float()\n",
    "        y_hat = model(x)\n",
    "        optimizer.zero_grad()\n",
    "        loss = loss_func(y_hat, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss_epoch += loss.item()\n",
    "        total_tmp += 1\n",
    "    auc = get_auc(test_loader, model)\n",
    "    print('epoch/epoches: {}/{}, train loss: {:.3f}, test auc: {:.3f}'.format(epoch, epoches, total_loss_epoch / total_tmp, auc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Criteo dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class CriteoDataset(torch.utils.data.Dataset):\n",
    "    \"\"\"\n",
    "    Criteo Display Advertising Challenge Dataset\n",
    "    Data prepration:\n",
    "        * Remove the infrequent features (appearing in less than threshold instances) and treat them as a single feature\n",
    "        * Discretize numerical values by log2 transformation which is proposed by the winner of Criteo Competition\n",
    "    :param dataset_path: criteo train.txt path.\n",
    "    :param cache_path: lmdb cache path.\n",
    "    :param rebuild_cache: If True, lmdb cache is refreshed.\n",
    "    :param min_threshold: infrequent feature threshold.\n",
    "    Reference:\n",
    "        https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset\n",
    "        https://www.csie.ntu.edu.tw/~r01922136/kaggle-2014-criteo.pdf\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dataset_path=None, cache_path='.criteo', rebuild_cache=False, min_threshold=10):\n",
    "        self.NUM_FEATS = 39\n",
    "        self.NUM_INT_FEATS = 13\n",
    "        self.min_threshold = min_threshold\n",
    "        if rebuild_cache or not Path(cache_path).exists():\n",
    "            shutil.rmtree(cache_path, ignore_errors=True)\n",
    "            if dataset_path is None:\n",
    "                raise ValueError('create cache: failed: dataset_path is None')\n",
    "            self.__build_cache(dataset_path, cache_path)\n",
    "        self.env = lmdb.open(cache_path, create=False, lock=False, readonly=True)\n",
    "        with self.env.begin(write=False) as txn:\n",
    "            self.length = txn.stat()['entries'] - 1\n",
    "            self.field_dims = np.frombuffer(txn.get(b'field_dims'), dtype=np.uint32)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        with self.env.begin(write=False) as txn:\n",
    "            np_array = np.frombuffer(\n",
    "                txn.get(struct.pack('>I', index)), dtype=np.uint32).astype(dtype=np.long)\n",
    "        return np_array[1:], np_array[0]\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "\n",
    "    def __build_cache(self, path, cache_path):\n",
    "        feat_mapper, defaults = self.__get_feat_mapper(path)\n",
    "        with lmdb.open(cache_path, map_size=int(1e11)) as env:\n",
    "            field_dims = np.zeros(self.NUM_FEATS, dtype=np.uint32)\n",
    "            for i, fm in feat_mapper.items():\n",
    "                field_dims[i - 1] = len(fm) + 1\n",
    "            with env.begin(write=True) as txn:\n",
    "                txn.put(b'field_dims', field_dims.tobytes())\n",
    "            for buffer in self.__yield_buffer(path, feat_mapper, defaults):\n",
    "                with env.begin(write=True) as txn:\n",
    "                    for key, value in buffer:\n",
    "                        txn.put(key, value)\n",
    "\n",
    "    def __get_feat_mapper(self, path):\n",
    "        feat_cnts = defaultdict(lambda: defaultdict(int))\n",
    "        with open(path) as f:\n",
    "            pbar = tqdm(f, mininterval=1, smoothing=0.1)\n",
    "            pbar.set_description('Create criteo dataset cache: counting features')\n",
    "            for line in pbar:\n",
    "                values = line.rstrip('\\n').split('\\t')\n",
    "                if len(values) != self.NUM_FEATS + 1:\n",
    "                    continue\n",
    "                for i in range(1, self.NUM_INT_FEATS + 1):\n",
    "                    feat_cnts[i][convert_numeric_feature(values[i])] += 1\n",
    "                for i in range(self.NUM_INT_FEATS + 1, self.NUM_FEATS + 1):\n",
    "                    feat_cnts[i][values[i]] += 1\n",
    "        feat_mapper = {i: {feat for feat, c in cnt.items() if c >= self.min_threshold} for i, cnt in feat_cnts.items()}\n",
    "        feat_mapper = {i: {feat: idx for idx, feat in enumerate(cnt)} for i, cnt in feat_mapper.items()}\n",
    "        defaults = {i: len(cnt) for i, cnt in feat_mapper.items()}\n",
    "        return feat_mapper, defaults\n",
    "\n",
    "    def __yield_buffer(self, path, feat_mapper, defaults, buffer_size=int(1e5)):\n",
    "        item_idx = 0\n",
    "        buffer = list()\n",
    "        with open(path) as f:\n",
    "            pbar = tqdm(f, mininterval=1, smoothing=0.1)\n",
    "            pbar.set_description('Create criteo dataset cache: setup lmdb')\n",
    "            for line in pbar:\n",
    "                values = line.rstrip('\\n').split('\\t')\n",
    "                if len(values) != self.NUM_FEATS + 1:\n",
    "                    continue\n",
    "                np_array = np.zeros(self.NUM_FEATS + 1, dtype=np.uint32)\n",
    "                np_array[0] = int(values[0])\n",
    "                for i in range(1, self.NUM_INT_FEATS + 1):\n",
    "                    np_array[i] = feat_mapper[i].get(convert_numeric_feature(values[i]), defaults[i])\n",
    "                for i in range(self.NUM_INT_FEATS + 1, self.NUM_FEATS + 1):\n",
    "                    np_array[i] = feat_mapper[i].get(values[i], defaults[i])\n",
    "                buffer.append((struct.pack('>I', item_idx), np_array.tobytes()))\n",
    "                item_idx += 1\n",
    "                if item_idx % buffer_size == 0:\n",
    "                    yield buffer\n",
    "                    buffer.clear()\n",
    "            yield buffer\n",
    "\n",
    "\n",
    "@lru_cache(maxsize=None)\n",
    "def convert_numeric_feature(val: str):\n",
    "    if val == '':\n",
    "        return 'NULL'\n",
    "    v = int(val)\n",
    "    if v > 2:\n",
    "        return str(int(math.log(v) ** 2))\n",
    "    else:\n",
    "        return str(v - 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/dataset/criteo.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2022-01-07 08:45:18\n",
      "\n",
      "recohut: 0.0.9\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.144+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "matplotlib: 3.2.2\n",
      "pandas    : 1.1.5\n",
      "lmdb      : 0.99\n",
      "PIL       : 7.1.2\n",
      "IPython   : 5.5.0\n",
      "numpy     : 1.19.5\n",
      "torch     : 1.10.0+cu111\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d -p recohut"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
