{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.dssm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DSSM\n",
    "> An implementation of DSSM, Deep Structured Semantic Model.\n",
    "\n",
    "Reference: https://github.com/massquantity/DBRL/blob/master/dbrl/models/dssm.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class DSSM(nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            main_embed_size,\n",
    "            feat_embed_size,\n",
    "            n_users,\n",
    "            n_items,\n",
    "            hidden_size,\n",
    "            feat_map,\n",
    "            static_feat,\n",
    "            dynamic_feat,\n",
    "            use_bn=True\n",
    "    ):\n",
    "        super(DSSM, self).__init__()\n",
    "        self.total_feat = static_feat + dynamic_feat\n",
    "        self.embed_user = nn.Embedding(n_users + 1, main_embed_size,\n",
    "                                       padding_idx=n_users)\n",
    "        self.embed_item = nn.Embedding(n_items + 1, main_embed_size,\n",
    "                                       padding_idx=n_items)\n",
    "        self.embed_feat = nn.ModuleDict({\n",
    "            feat: nn.Embedding(feat_map[feat + \"_vocab\"] + 1, feat_embed_size,\n",
    "                               padding_idx=feat_map[feat + \"_vocab\"])\n",
    "            for feat in self.total_feat\n",
    "        })\n",
    "\n",
    "        self.static_feat = static_feat\n",
    "        self.dynamic_feat = dynamic_feat\n",
    "        input_dim_user = main_embed_size + feat_embed_size * len(static_feat)\n",
    "        self.fcu1 = nn.Linear(input_dim_user, hidden_size[0])\n",
    "        self.fcu2 = nn.Linear(hidden_size[0], hidden_size[1])\n",
    "        self.fcu3 = nn.Linear(hidden_size[1], main_embed_size)\n",
    "\n",
    "        input_dim_item = main_embed_size + feat_embed_size * len(dynamic_feat)\n",
    "        self.fci1 = nn.Linear(input_dim_item, hidden_size[0])\n",
    "        self.fci2 = nn.Linear(hidden_size[0], hidden_size[1])\n",
    "        self.fci3 = nn.Linear(hidden_size[1], main_embed_size)\n",
    "\n",
    "        self.use_bn = use_bn\n",
    "        if use_bn:\n",
    "            self.bnu1 = nn.BatchNorm1d(hidden_size[0])\n",
    "            self.bnu2 = nn.BatchNorm1d(hidden_size[1])\n",
    "            self.bni1 = nn.BatchNorm1d(hidden_size[0])\n",
    "            self.bni2 = nn.BatchNorm1d(hidden_size[1])\n",
    "\n",
    "    def get_embedding(self, data):\n",
    "        user_part = [self.embed_user(data[\"user\"])]\n",
    "        for feat in self.static_feat:\n",
    "            embedding = self.embed_feat[feat]\n",
    "            user_part.append(embedding(data[feat]))\n",
    "\n",
    "        user_part = torch.cat(user_part, dim=1)\n",
    "        out_user = self.fcu1(user_part)\n",
    "        if self.use_bn:\n",
    "            out_user = self.bnu1(out_user)\n",
    "        out_user = F.relu(out_user)\n",
    "        out_user = self.fcu2(out_user)\n",
    "        if self.use_bn:\n",
    "            out_user = self.bnu2(out_user)\n",
    "        out_user = F.relu(out_user)\n",
    "        out_user = self.fcu3(out_user)\n",
    "        out_user = out_user / torch.norm(out_user, dim=1, keepdim=True)\n",
    "\n",
    "        item_part = [self.embed_item(data[\"item\"])]\n",
    "        for feat in self.dynamic_feat:\n",
    "            embedding = self.embed_feat[feat]\n",
    "            item_part.append(embedding(data[feat]))\n",
    "\n",
    "        item_part = torch.cat(item_part, dim=1)\n",
    "        out_item = self.fci1(item_part)\n",
    "        if self.use_bn:\n",
    "            out_item = self.bni1(out_item)\n",
    "        out_item = F.relu(out_item)\n",
    "        out_item = self.fci2(out_item)\n",
    "        if self.use_bn:\n",
    "            out_item = self.bni2(out_item)\n",
    "        out_item = F.relu(out_item)\n",
    "        out_item = self.fci3(out_item)\n",
    "        out_item = out_item / torch.norm(out_item, dim=1, keepdim=True)\n",
    "        return out_user, out_item\n",
    "\n",
    "    def forward(self, data):\n",
    "        out_user, out_item = self.get_embedding(data)\n",
    "        out = torch.sum(torch.mul(out_user, out_item), dim=1).squeeze()\n",
    "        return out_user, out_item, out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DSSM(\n",
       "  (embed_user): Embedding(6, 2, padding_idx=5)\n",
       "  (embed_item): Embedding(6, 2, padding_idx=5)\n",
       "  (embed_feat): ModuleDict(\n",
       "    (a): Embedding(3, 2, padding_idx=2)\n",
       "    (b): Embedding(3, 2, padding_idx=2)\n",
       "  )\n",
       "  (fcu1): Linear(in_features=4, out_features=5, bias=True)\n",
       "  (fcu2): Linear(in_features=5, out_features=3, bias=True)\n",
       "  (fcu3): Linear(in_features=3, out_features=2, bias=True)\n",
       "  (fci1): Linear(in_features=4, out_features=5, bias=True)\n",
       "  (fci2): Linear(in_features=5, out_features=3, bias=True)\n",
       "  (fci3): Linear(in_features=3, out_features=2, bias=True)\n",
       "  (bnu1): BatchNorm1d(5, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (bnu2): BatchNorm1d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (bni1): BatchNorm1d(5, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (bni2): BatchNorm1d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = DSSM(main_embed_size=2, feat_embed_size=2, n_users=5, n_items=5, hidden_size=[5,3], feat_map={'a_vocab':2,'b_vocab':2}, static_feat=['a'], dynamic_feat=['b'], use_bn=True)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-24 04:48:03\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",
      "IPython: 5.5.0\n",
      "torch  : 1.10.0+cu111\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
