{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp datasets.avazu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Avazu\n",
    "> Avazu 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 shutil\n",
    "import struct\n",
    "from collections import defaultdict\n",
    "from pathlib import Path\n",
    "\n",
    "import lmdb\n",
    "import numpy as np\n",
    "import torch.utils.data\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class AvazuDataset(torch.utils.data.Dataset):\n",
    "    \"\"\"\n",
    "    Avazu Click-Through Rate Prediction Dataset\n",
    "    Dataset preparation\n",
    "        Remove the infrequent features (appearing in less than threshold instances) and treat them as a single feature\n",
    "    :param dataset_path: avazu train 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://www.kaggle.com/c/avazu-ctr-prediction\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dataset_path=None, cache_path='.avazu', rebuild_cache=False, min_threshold=4):\n",
    "        self.NUM_FEATS = 22\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",
    "            f.readline()\n",
    "            pbar = tqdm(f, mininterval=1, smoothing=0.1)\n",
    "            pbar.set_description('Create avazu dataset cache: counting features')\n",
    "            for line in pbar:\n",
    "                values = line.rstrip('\\n').split(',')\n",
    "                if len(values) != self.NUM_FEATS + 2:\n",
    "                    continue\n",
    "                for i in range(1, self.NUM_FEATS + 1):\n",
    "                    feat_cnts[i][values[i + 1]] += 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",
    "            f.readline()\n",
    "            pbar = tqdm(f, mininterval=1, smoothing=0.1)\n",
    "            pbar.set_description('Create avazu dataset cache: setup lmdb')\n",
    "            for line in pbar:\n",
    "                values = line.rstrip('\\n').split(',')\n",
    "                if len(values) != self.NUM_FEATS + 2:\n",
    "                    continue\n",
    "                np_array = np.zeros(self.NUM_FEATS + 1, dtype=np.uint32)\n",
    "                np_array[0] = int(values[1])\n",
    "                for i in range(1, self.NUM_FEATS + 1):\n",
    "                    np_array[i] = feat_mapper[i].get(values[i+1], 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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/dataset/avazu.py"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
