notebooks/two_tower_model_clothing_recommendation_system.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3fcd11be-9d88-4c0c-9d26-6f7d33dbd311",
   "metadata": {},
   "source": [
    "# Clothing recommendation system based on two tower model"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "This notebook contains the implementation of a recommendation system based on [two tower model architecture](https://research.google/pubs/pub48840/) and inspections related to it. The two tower model is implemented, trained and tested with [PyTorch framework](https://pytorch.org/). To simplify operations on the structured dataset, [pandas](https://pandas.pydata.org/) library is used. The used dataset for the inspections of the model is the [H&M Personalized Fashion Recommendations dataset](https://www.kaggle.com/competitions/h-and-m-personalized-fashion-recommendations/). To track the experiments, run hyperparameter search and save model artifacts [Weights & Biases (wandb)](https://wandb.ai/) is used. The model is exported and saved in [ONNX](https://onnx.ai/) format. The notebook is written in Python."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6e071d5b10973b92"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 1. Two tower model implementation and experimentation"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "292120fc9847f81e"
  },
  {
   "cell_type": "markdown",
   "source": [
    "Let's import the required dependencies first! If a dependency can't be imported, because it is missing from the system, then [pip](https://pypi.org/project/pip/) can be used for instance to download them."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5d0c5e91a5906743"
  },
  {
   "cell_type": "code",
   "source": [
    "import bisect\n",
    "import json\n",
    "import math\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import wandb\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torch import nn\n",
    "import traceback"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6e734d62e5b72a50",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "To make steps reproducible, that contain random behaviour, let's set the random generator's seed to a constant number."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "a454cb55639529c5"
  },
  {
   "cell_type": "code",
   "source": [
    "np.random.seed(2)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "bcdf132caf91d1c7",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "To achieve different goals, different datasets are can be more suited. But datasets that are use the data from the H&M fashion dataset might require similar functionalities, like attribute transformations. The **WebshopDataContainer** class contains these common implementations, that the more specific datasets might reuse."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "380084d85b24e85d"
  },
  {
   "cell_type": "code",
   "source": [
    "class WebshopDataContainer:\n",
    "    _ARTICLE_GARMENT_GROUP_NAME_NUMERIC_ENCODER = {\n",
    "        \"Accessories\": 0,\n",
    "        \"Blouses\": 1,\n",
    "        \"Dressed\": 2,\n",
    "        \"Dresses Ladies\": 3,\n",
    "        \"Dresses/Skirts girls\": 4,\n",
    "        \"Jersey Basic\": 5,\n",
    "        \"Jersey Fancy\": 6,\n",
    "        \"Knitwear\": 7,\n",
    "        \"Outdoor\": 8,\n",
    "        \"Shirts\": 9,\n",
    "        \"Shoes\": 10,\n",
    "        \"Shorts\": 11,\n",
    "        \"Skirts\": 12,\n",
    "        \"Socks and Tights\": 13,\n",
    "        \"Special Offers\": 14,\n",
    "        \"Swimwear\": 15,\n",
    "        \"Trousers\": 16,\n",
    "        \"Trousers Denim\": 17,\n",
    "        \"Under-, Nightwear\": 18,\n",
    "        \"Unknown\": 19,\n",
    "        \"Woven/Jersey/Knitted mix Baby\": 20,\n",
    "    }\n",
    "\n",
    "    _ARTICLE_GRAPHICAL_APPEARANCE_NUMERIC_ENCODER = {\n",
    "        \"All over pattern\": 0,\n",
    "        \"Application/3D\": 1,\n",
    "        \"Argyle\": 2,\n",
    "        \"Chambray\": 3,\n",
    "        \"Check\": 4,\n",
    "        \"Colour blocking\": 5,\n",
    "        \"Contrast\": 6,\n",
    "        \"Denim\": 7,\n",
    "        \"Dot\": 8,\n",
    "        \"Embroidery\": 9,\n",
    "        \"Front print\": 10,\n",
    "        \"Glittering/Metallic\": 11,\n",
    "        \"Hologram\": 12,\n",
    "        \"Jacquard\": 13,\n",
    "        \"Lace\": 14,\n",
    "        \"Melange\": 15,\n",
    "        \"Mesh\": 16,\n",
    "        \"Metallic\": 17,\n",
    "        \"Mixed solid/pattern\": 18,\n",
    "        \"Neps\": 19,\n",
    "        \"Other pattern\": 20,\n",
    "        \"Other structure\": 21,\n",
    "        \"Placement print\": 22,\n",
    "        \"Sequin\": 23,\n",
    "        \"Slub\": 24,\n",
    "        \"Solid\": 25,\n",
    "        \"Stripe\": 26,\n",
    "        \"Transparent\": 27,\n",
    "        \"Treatment\": 28,\n",
    "        \"Unknown\": 29,\n",
    "    }\n",
    "\n",
    "    _ARTICLE_INDEX_NAME_NUMERIC_ENCODER = {\n",
    "        \"Baby Sizes 50-98\": 0,\n",
    "        \"Children Accessories, Swimwear\": 1,\n",
    "        \"Children Sizes 134-170\": 2,\n",
    "        \"Children Sizes 92-140\": 3,\n",
    "        \"Divided\": 4,\n",
    "        \"Ladies Accessories\": 5,\n",
    "        \"Ladieswear\": 6,\n",
    "        \"Lingeries/Tights\": 7,\n",
    "        \"Menswear\": 8,\n",
    "        \"Sport\": 9,\n",
    "    }\n",
    "\n",
    "    _ARTICLE_PERCEIVED_COLOR_MASTER_NAME_NUMERIC_ENCODER = {\n",
    "        \"Beige\": 0,\n",
    "        \"Black\": 1,\n",
    "        \"Blue\": 2,\n",
    "        \"Bluish Green\": 3,\n",
    "        \"Brown\": 4,\n",
    "        \"Green\": 5,\n",
    "        \"Grey\": 6,\n",
    "        \"Khaki green\": 7,\n",
    "        \"Lilac Purple\": 8,\n",
    "        \"Metal\": 9,\n",
    "        \"Mole\": 10,\n",
    "        \"Orange\": 11,\n",
    "        \"Pink\": 12,\n",
    "        \"Red\": 13,\n",
    "        \"Turquoise\": 14,\n",
    "        \"undefined\": 15,\n",
    "        \"Unknown\": 16,\n",
    "        \"White\": 17,\n",
    "        \"Yellow\": 18,\n",
    "        \"Yellowish Green\": 19,\n",
    "    }\n",
    "\n",
    "    _ARTICLE_PRODUCT_GROUP_NAME_NUMERIC_ENCODER = {\n",
    "        \"Accessories\": 0,\n",
    "        \"Bags\": 1,\n",
    "        \"Cosmetic\": 2,\n",
    "        \"Fun\": 3,\n",
    "        \"Furniture\": 4,\n",
    "        \"Garment Full body\": 5,\n",
    "        \"Garment Lower body\": 6,\n",
    "        \"Garment Upper body\": 7,\n",
    "        \"Garment and Shoe care\": 8,\n",
    "        \"Interior textile\": 9,\n",
    "        \"Items\": 10,\n",
    "        \"Nightwear\": 11,\n",
    "        \"Shoes\": 12,\n",
    "        \"Socks & Tights\": 13,\n",
    "        \"Stationery\": 14,\n",
    "        \"Swimwear\": 15,\n",
    "        \"Underwear\": 16,\n",
    "        \"Underwear/nightwear\": 17,\n",
    "        \"Unknown\": 18,\n",
    "    }\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            articles_df,\n",
    "            customers_df,\n",
    "            transactions_df,\n",
    "            customer_age_group_size,\n",
    "            customer_minimum_age,\n",
    "            article_id_column_name,\n",
    "            article_garment_group_column_name,\n",
    "            article_graphical_appearance_column_name,\n",
    "            article_index_name_column_name,\n",
    "            article_perceived_color_master_column_name,\n",
    "            article_product_group_column_name,\n",
    "            customer_id_column_name,\n",
    "            customer_age_column_name,\n",
    "            transaction_date_column_name,\n",
    "    ):\n",
    "\n",
    "        self._articles_df = articles_df\n",
    "        self._customers_df = customers_df\n",
    "        self._transactions_df = transactions_df\n",
    "\n",
    "        self._article_id_column_name = article_id_column_name\n",
    "        self._article_garment_group_column_name = article_garment_group_column_name\n",
    "        self._article_graphical_appearance_column_name = article_graphical_appearance_column_name\n",
    "        self._article_index_name_column_name = article_index_name_column_name\n",
    "        self._article_perceived_color_master_column_name = article_perceived_color_master_column_name\n",
    "        self._article_product_group_column_name = article_product_group_column_name\n",
    "        self._customer_id_column_name = customer_id_column_name\n",
    "        self._customer_age_column_name = customer_age_column_name\n",
    "        self._transaction_date_column_name = transaction_date_column_name\n",
    "\n",
    "        self._customer_age_group_size = customer_age_group_size\n",
    "        self._customer_minimum_age = customer_minimum_age\n",
    "\n",
    "        if transaction_date_column_name in transactions_df.columns:\n",
    "            self._is_transaction_date_column_used = True\n",
    "            self._transaction_dates = transactions_df[transaction_date_column_name].unique()\n",
    "        else:\n",
    "            self._is_transaction_date_column_used = False\n",
    "\n",
    "    def _calculate_customer_age_group_index(self, age):\n",
    "        return math.ceil((age - self._customer_minimum_age) / self._customer_age_group_size)\n",
    "\n",
    "    def _create_result_attributes(self, article, customer, transaction):\n",
    "        article_id = article[self._article_id_column_name]\n",
    "        customer_id = customer[self._customer_id_column_name]\n",
    "        result_attributes = {\n",
    "            self._article_id_column_name: self._articles_df.index.get_loc(\n",
    "                self._articles_df[self._articles_df[self._article_id_column_name] == article_id].index[0]),\n",
    "            self._customer_id_column_name: self._customers_df.index.get_loc(\n",
    "                self._customers_df[self._customers_df[self._customer_id_column_name] == customer_id].index[0]),\n",
    "        }\n",
    "\n",
    "        if self._article_garment_group_column_name in self._articles_df.columns:\n",
    "            result_attributes[self._article_garment_group_column_name] = \\\n",
    "                self._ARTICLE_GARMENT_GROUP_NAME_NUMERIC_ENCODER[\n",
    "                    article[self._article_garment_group_column_name]\n",
    "                ]\n",
    "\n",
    "        if self._article_graphical_appearance_column_name in self._articles_df.columns:\n",
    "            result_attributes[self._article_graphical_appearance_column_name] = (\n",
    "                self._ARTICLE_GRAPHICAL_APPEARANCE_NUMERIC_ENCODER[\n",
    "                    article[self._article_graphical_appearance_column_name]\n",
    "                ]\n",
    "            )\n",
    "\n",
    "        if self._article_index_name_column_name in self._articles_df.columns:\n",
    "            result_attributes[self._article_index_name_column_name] = self._ARTICLE_INDEX_NAME_NUMERIC_ENCODER[\n",
    "                article[self._article_index_name_column_name]\n",
    "            ]\n",
    "\n",
    "        if self._article_perceived_color_master_column_name in self._articles_df.columns:\n",
    "            result_attributes[self._article_perceived_color_master_column_name] = (\n",
    "                self._ARTICLE_PERCEIVED_COLOR_MASTER_NAME_NUMERIC_ENCODER[\n",
    "                    article[self._article_perceived_color_master_column_name]\n",
    "                ])\n",
    "\n",
    "        if self._article_product_group_column_name in self._articles_df.columns:\n",
    "            result_attributes[self._article_product_group_column_name] = \\\n",
    "                self._ARTICLE_PRODUCT_GROUP_NAME_NUMERIC_ENCODER[\n",
    "                    article[self._article_product_group_column_name]\n",
    "                ]\n",
    "\n",
    "        if self._customer_age_column_name in self._customers_df.columns:\n",
    "            result_attributes[self._customer_age_column_name] = self._calculate_customer_age_group_index(\n",
    "                customer[self._customer_age_column_name]\n",
    "            )\n",
    "\n",
    "        if self._is_transaction_date_column_used:\n",
    "            result_attributes[self._transaction_date_column_name] = np.where(\n",
    "                self._transaction_dates == transaction[self._transaction_date_column_name]\n",
    "            )[0][0]\n",
    "\n",
    "        return result_attributes\n",
    "\n",
    "    @classmethod\n",
    "    def get_article_garment_group_name_numeric_encoder(cls):\n",
    "        return cls._ARTICLE_GARMENT_GROUP_NAME_NUMERIC_ENCODER.copy()\n",
    "\n",
    "    @classmethod\n",
    "    def get_article_graphical_appearance_numeric_encoder(cls):\n",
    "        return cls._ARTICLE_GRAPHICAL_APPEARANCE_NUMERIC_ENCODER.copy()\n",
    "\n",
    "    @classmethod\n",
    "    def get_article_index_name_numeric_encoder(cls):\n",
    "        return cls._ARTICLE_INDEX_NAME_NUMERIC_ENCODER.copy()\n",
    "\n",
    "    @classmethod\n",
    "    def get_article_perceived_color_master_name_numeric_encoder(cls):\n",
    "        return cls._ARTICLE_PERCEIVED_COLOR_MASTER_NAME_NUMERIC_ENCODER.copy()\n",
    "\n",
    "    @classmethod\n",
    "    def get_article_product_group_name_numeric_encoder(cls):\n",
    "        return cls._ARTICLE_PRODUCT_GROUP_NAME_NUMERIC_ENCODER.copy()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "93291b10004b7b35",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The **NegativeBatchSampledTransactionDataset** stores the data of the transactions. The speciality of it (as its name suggests) that its also generates negative transaction samples. For every real transaction a negative transaction with the same date and customer generated and its article is chosen in such a way, that the customer has no real transaction with that. The information, that the transaction is real or a negative example, is stored as the label value."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "992902344d7e6137"
  },
  {
   "cell_type": "code",
   "source": [
    "class NegativeBatchSampledTransactionDataset(Dataset, WebshopDataContainer):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            articles_df,\n",
    "            customers_df,\n",
    "            transactions_df,\n",
    "            customer_age_group_size,\n",
    "            customer_minimum_age,\n",
    "            article_id_column_name,\n",
    "            article_garment_group_column_name,\n",
    "            article_graphical_appearance_column_name,\n",
    "            article_index_name_column_name,\n",
    "            article_perceived_color_master_column_name,\n",
    "            article_product_group_column_name,\n",
    "            customer_id_column_name,\n",
    "            customer_age_column_name,\n",
    "            transaction_date_column_name,\n",
    "            get_negative_transactions_from_cache: bool = False,\n",
    "    ):\n",
    "        WebshopDataContainer.__init__(\n",
    "            self=self,\n",
    "            articles_df=articles_df,\n",
    "            customers_df=customers_df,\n",
    "            transactions_df=transactions_df,\n",
    "            customer_age_group_size=customer_age_group_size,\n",
    "            customer_minimum_age=customer_minimum_age,\n",
    "            article_id_column_name=article_id_column_name,\n",
    "            article_garment_group_column_name=article_garment_group_column_name,\n",
    "            article_graphical_appearance_column_name=article_graphical_appearance_column_name,\n",
    "            article_index_name_column_name=article_index_name_column_name,\n",
    "            article_perceived_color_master_column_name=article_perceived_color_master_column_name,\n",
    "            article_product_group_column_name=article_product_group_column_name,\n",
    "            customer_id_column_name=customer_id_column_name,\n",
    "            customer_age_column_name=customer_age_column_name,\n",
    "            transaction_date_column_name=transaction_date_column_name,\n",
    "        )\n",
    "\n",
    "        self._cached_negative_transactions_path \\\n",
    "            = \"two_tower_model_clothing_recommendation_system_cache/cached_negative_transactions.json\"\n",
    "\n",
    "        if get_negative_transactions_from_cache:\n",
    "            try:\n",
    "                with open(self._cached_negative_transactions_path, \"r\") as json_file:\n",
    "                    self._negative_article_ids_of_customers = json.load(json_file)\n",
    "            except FileNotFoundError:\n",
    "                self._negative_article_ids_of_customers = None\n",
    "        else:\n",
    "            self._negative_article_ids_of_customers = None\n",
    "\n",
    "        if self._negative_article_ids_of_customers is None:\n",
    "            self._negative_article_ids_of_customers = self._create_negative_article_ids_of_customers(\n",
    "                articles_df=self._articles_df,\n",
    "                transactions_df=self._transactions_df,\n",
    "            )\n",
    "            try:\n",
    "                with open(self._cached_negative_transactions_path, \"w\") as json_file:\n",
    "                    json.dump(self._negative_article_ids_of_customers, json_file)\n",
    "            except (PermissionError, FileNotFoundError, OSError) as e:\n",
    "                print(f\"Failed to save negative transactions! {e}\")\n",
    "\n",
    "        self.transaction_counter_for_customer = {}\n",
    "\n",
    "    def _create_negative_article_ids_of_customers(self, articles_df, transactions_df):\n",
    "        negative_article_ids_of_customers = {}\n",
    "        customer_count = transactions_df[self._customer_id_column_name].nunique()\n",
    "        print(\"Starting to generate negative articles for customers...\")\n",
    "\n",
    "        for index, customer_id in enumerate(transactions_df[self._customer_id_column_name].drop_duplicates()):\n",
    "            transactions_of_customer = transactions_df[transactions_df[self._customer_id_column_name] == customer_id]\n",
    "            ids_of_articles_purchased_by_customer = transactions_of_customer[self._article_id_column_name].drop_duplicates()\n",
    "\n",
    "            ids_of_negative_articles_of_customer = (\n",
    "                articles_df[~articles_df[self._article_id_column_name].isin(ids_of_articles_purchased_by_customer)]\n",
    "                [self._article_id_column_name]\n",
    "            )\n",
    "\n",
    "            negative_article_ids_of_customers[customer_id] = np.random.choice(\n",
    "                ids_of_negative_articles_of_customer,\n",
    "                size=len(transactions_of_customer),\n",
    "                replace=False,\n",
    "            ).tolist()\n",
    "\n",
    "            if index % 1000 == 0:\n",
    "                print(f\"Negative article generation progress: {index}/{customer_count}\")\n",
    "\n",
    "        print(\"Finished generating negative articles for customers!\")\n",
    "        return negative_article_ids_of_customers\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self._transactions_df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        transaction = self._transactions_df.iloc[idx]\n",
    "        customer = self._customers_df[\n",
    "            self._customers_df[self._customer_id_column_name] == transaction[self._customer_id_column_name]\n",
    "            ].iloc[0]\n",
    "        positive_article = self._articles_df[\n",
    "            self._articles_df[self._article_id_column_name] == transaction[self._article_id_column_name]\n",
    "            ].iloc[0]\n",
    "\n",
    "        customer_id = transaction[self._customer_id_column_name]\n",
    "        transaction_counter_of_customer = self.transaction_counter_for_customer.get(customer_id, 0)\n",
    "        negative_article_id = self._negative_article_ids_of_customers[customer_id][transaction_counter_of_customer]\n",
    "        negative_article = (\n",
    "            self._articles_df[self._articles_df[self._article_id_column_name] == negative_article_id].iloc[0]\n",
    "        )\n",
    "\n",
    "        negative_transaction = pd.Series({\n",
    "            self._article_id_column_name: negative_article_id,\n",
    "            self._customer_id_column_name: customer_id,\n",
    "        })\n",
    "\n",
    "        if self._transaction_date_column_name in self._transactions_df.columns:\n",
    "            negative_transaction[self._transaction_date_column_name] = transaction[self._transaction_date_column_name]\n",
    "\n",
    "        self.transaction_counter_for_customer[customer_id] = transaction_counter_of_customer + 1\n",
    "\n",
    "        positive_transaction_attributes = self._create_result_attributes(\n",
    "            article=positive_article,\n",
    "            customer=customer,\n",
    "            transaction=transaction,\n",
    "        )\n",
    "        negative_transactions_attributes = self._create_result_attributes(\n",
    "            article=negative_article,\n",
    "            customer=customer,\n",
    "            transaction=negative_transaction,\n",
    "        )\n",
    "\n",
    "        return (\n",
    "            [positive_transaction_attributes, negative_transactions_attributes],\n",
    "            [1.0, 0.0],\n",
    "        )"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2f956a20f3dfafdd",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The **TransactionDataset** is a simple dataset. It just maps the attributes of transactions and related entities (article, customer), and returns them."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "d7834baa162b9cd8"
  },
  {
   "cell_type": "code",
   "source": [
    "class TransactionDataset(Dataset, WebshopDataContainer):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            articles_df,\n",
    "            customers_df,\n",
    "            transactions_df,\n",
    "            customer_age_group_size,\n",
    "            customer_minimum_age,\n",
    "            article_id_column_name,\n",
    "            article_garment_group_column_name,\n",
    "            article_graphical_appearance_column_name,\n",
    "            article_index_name_column_name,\n",
    "            article_perceived_color_master_column_name,\n",
    "            article_product_group_column_name,\n",
    "            customer_id_column_name,\n",
    "            customer_age_column_name,\n",
    "            transaction_date_column_name,\n",
    "    ):\n",
    "        WebshopDataContainer.__init__(\n",
    "            self=self,\n",
    "            articles_df=articles_df,\n",
    "            customers_df=customers_df,\n",
    "            transactions_df=transactions_df,\n",
    "            customer_age_group_size=customer_age_group_size,\n",
    "            customer_minimum_age=customer_minimum_age,\n",
    "            article_id_column_name=article_id_column_name,\n",
    "            article_garment_group_column_name=article_garment_group_column_name,\n",
    "            article_graphical_appearance_column_name=article_graphical_appearance_column_name,\n",
    "            article_index_name_column_name=article_index_name_column_name,\n",
    "            article_perceived_color_master_column_name=article_perceived_color_master_column_name,\n",
    "            article_product_group_column_name=article_product_group_column_name,\n",
    "            customer_id_column_name=customer_id_column_name,\n",
    "            customer_age_column_name=customer_age_column_name,\n",
    "            transaction_date_column_name=transaction_date_column_name,\n",
    "        )\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self._transactions_df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        transaction = self._transactions_df.iloc[idx]\n",
    "        customer = self._customers_df[\n",
    "            self._customers_df[self._customer_id_column_name] == transaction[self._customer_id_column_name]\n",
    "            ].iloc[0]\n",
    "        positive_article = self._articles_df[\n",
    "            self._articles_df[self._article_id_column_name] == transaction[self._article_id_column_name]\n",
    "            ].iloc[0]\n",
    "\n",
    "        transaction_attributes = self._create_result_attributes(\n",
    "            article=positive_article,\n",
    "            customer=customer,\n",
    "            transaction=transaction,\n",
    "        )\n",
    "\n",
    "        return transaction_attributes, 1.0"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "57ea6b8e4bd0369f",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The **TwoTowerModel** class contains the implementation of the recommender model. The tower model contains two separate models, usually both model is an identical neural network. Each tower model receives the input data of one of the related entity type. In this case the item tower calculates the output using the data of the article entities, the query tower uses the data of the customer entities. The tower models calculate an embedding vector, and the goal of the model is to create similar vectors of such entities that are connected. In the current example the connection between articles and customers is whether they appear together in a transaction. Dense categorical attributes (typically enum values) are encoded with one hot encoding, sparse categorical attributes are encoded using embedding layers. The embedding layers' dictionaries' size are incremented with one, because if the recommendation system is deployed, it is possible, that it will get a value, that was not contained in its test dataset. These unknown values should be passed as this special highest index value as input."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6edc716a8ef34791"
  },
  {
   "cell_type": "code",
   "source": [
    "class TwoTowerModel(nn.Module):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            device,\n",
    "            number_of_neurons_in_the_first_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_second_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_third_hidden_dense_layer,\n",
    "            article_id_embedding_dimension_and_unique_value_count,\n",
    "            article_garment_group_name_unique_value_count,\n",
    "            article_graphical_appearance_name_unique_value_count,\n",
    "            article_index_name_unique_value_count,\n",
    "            article_perceived_color_master_name_unique_value_count,\n",
    "            article_product_group_name_unique_value_count,\n",
    "            customer_id_embedding_dimension_and_unique_value_count,\n",
    "            customer_age_group_count,\n",
    "            transaction_date_embedding_dimension_and_unique_value_count,\n",
    "            similarity_function,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        \n",
    "        self._device = device\n",
    "\n",
    "        self._initialize_input_layer(\n",
    "            article_id_embedding_dimension_and_unique_value_count=article_id_embedding_dimension_and_unique_value_count,\n",
    "            article_index_name_unique_value_count=article_index_name_unique_value_count,\n",
    "            article_garment_group_name_unique_value_count=article_garment_group_name_unique_value_count,\n",
    "            article_product_group_name_unique_value_count=article_product_group_name_unique_value_count,\n",
    "            article_perceived_color_master_name_unique_value_count=article_perceived_color_master_name_unique_value_count,\n",
    "            article_graphical_appearance_name_unique_value_count=article_graphical_appearance_name_unique_value_count,\n",
    "            customer_id_embedding_dimension_and_unique_value_count=customer_id_embedding_dimension_and_unique_value_count,\n",
    "            customer_age_group_count=customer_age_group_count,\n",
    "            transaction_date_embedding_dimension_and_unique_value_count=transaction_date_embedding_dimension_and_unique_value_count,\n",
    "        )\n",
    "\n",
    "        self._create_item_tower_model(\n",
    "            dimension_of_input_for_first_dense_layer=self.calculate_item_tower_input_dimension(),\n",
    "            number_of_neurons_in_the_first_hidden_dense_layer=number_of_neurons_in_the_first_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_second_hidden_dense_layer=number_of_neurons_in_the_second_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_third_hidden_dense_layer=number_of_neurons_in_the_third_hidden_dense_layer,\n",
    "        )\n",
    "        self._create_query_tower_model(\n",
    "            dimension_of_input_for_first_dense_layer=self.calculate_query_tower_input_dimension(),\n",
    "            number_of_neurons_in_the_first_hidden_dense_layer=number_of_neurons_in_the_first_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_second_hidden_dense_layer=number_of_neurons_in_the_second_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_third_hidden_dense_layer=number_of_neurons_in_the_third_hidden_dense_layer,\n",
    "        )\n",
    "\n",
    "        self.similarity_function = similarity_function\n",
    "\n",
    "    def _initialize_input_layer(\n",
    "            self,\n",
    "            article_id_embedding_dimension_and_unique_value_count,\n",
    "            article_index_name_unique_value_count,\n",
    "            article_garment_group_name_unique_value_count,\n",
    "            article_product_group_name_unique_value_count,\n",
    "            article_perceived_color_master_name_unique_value_count,\n",
    "            article_graphical_appearance_name_unique_value_count,\n",
    "            customer_id_embedding_dimension_and_unique_value_count,\n",
    "            customer_age_group_count,\n",
    "            transaction_date_embedding_dimension_and_unique_value_count,\n",
    "    ):\n",
    "        # TODO prepare one-hot encoded attribute processing for unknown values\n",
    "        if article_id_embedding_dimension_and_unique_value_count is not None:\n",
    "            self._article_id_embedding_dimension = article_id_embedding_dimension_and_unique_value_count[0]\n",
    "            self.article_id_encoder = nn.Embedding(\n",
    "                num_embeddings=article_id_embedding_dimension_and_unique_value_count[1] + 1,\n",
    "                embedding_dim=article_id_embedding_dimension_and_unique_value_count[0],\n",
    "            )\n",
    "        else:\n",
    "            self._article_id_embedding_dimension = 0\n",
    "\n",
    "        if customer_id_embedding_dimension_and_unique_value_count is not None:\n",
    "            self._customer_id_embedding_dimension = customer_id_embedding_dimension_and_unique_value_count[0]\n",
    "            self.customer_id_encoder = nn.Embedding(\n",
    "                num_embeddings=customer_id_embedding_dimension_and_unique_value_count[1] + 1,\n",
    "                embedding_dim=customer_id_embedding_dimension_and_unique_value_count[0],\n",
    "            )\n",
    "        else:\n",
    "            self._customer_id_embedding_dimension = 0\n",
    "\n",
    "        if transaction_date_embedding_dimension_and_unique_value_count is not None \\\n",
    "                and transaction_date_embedding_dimension_and_unique_value_count != 0:\n",
    "            self._customer_transaction_date_dimension = transaction_date_embedding_dimension_and_unique_value_count[0]\n",
    "            self.transaction_date_encoder = nn.Embedding(\n",
    "                num_embeddings=transaction_date_embedding_dimension_and_unique_value_count[1] + 1,\n",
    "                embedding_dim=transaction_date_embedding_dimension_and_unique_value_count[0],\n",
    "            )\n",
    "        else:\n",
    "            self._customer_transaction_date_dimension = 0\n",
    "            self.transaction_date_encoder = None\n",
    "\n",
    "        self._customer_age_group_unique_value_count = customer_age_group_count\n",
    "\n",
    "        self._article_index_name_unique_value_count = article_index_name_unique_value_count\n",
    "        self._article_garment_group_name_unique_value_count = article_garment_group_name_unique_value_count\n",
    "        self._article_product_group_name_unique_value_count = article_product_group_name_unique_value_count\n",
    "        self._article_perceived_color_master_name_unique_value_count = article_perceived_color_master_name_unique_value_count\n",
    "        self._article_graphical_appearance_name_unique_value_count = article_graphical_appearance_name_unique_value_count\n",
    "\n",
    "    def _create_item_tower_model(\n",
    "            self,\n",
    "            dimension_of_input_for_first_dense_layer,\n",
    "            number_of_neurons_in_the_first_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_second_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_third_hidden_dense_layer,\n",
    "    ):\n",
    "        self.item_tower_model = self._create_dense_sequential_stack(\n",
    "            list_of_layer_input_and_output_dimensions=self._assemble_list_of_sequential_dense_linear_stack_dimensions(\n",
    "                dimension_of_input_for_first_dense_layer=dimension_of_input_for_first_dense_layer,\n",
    "                number_of_neurons_in_the_first_hidden_dense_layer=number_of_neurons_in_the_first_hidden_dense_layer,\n",
    "                number_of_neurons_in_the_second_hidden_dense_layer=number_of_neurons_in_the_second_hidden_dense_layer,\n",
    "                number_of_neurons_in_the_third_hidden_dense_layer=number_of_neurons_in_the_third_hidden_dense_layer,\n",
    "            ),\n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def _assemble_list_of_sequential_dense_linear_stack_dimensions(\n",
    "            cls,\n",
    "            dimension_of_input_for_first_dense_layer,\n",
    "            number_of_neurons_in_the_first_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_second_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_third_hidden_dense_layer,\n",
    "    ):\n",
    "        list_of_tower_layer_input_and_output_dimensions = [\n",
    "            dimension_of_input_for_first_dense_layer,\n",
    "            number_of_neurons_in_the_first_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_second_hidden_dense_layer,\n",
    "        ]\n",
    "\n",
    "        if number_of_neurons_in_the_third_hidden_dense_layer is not None or number_of_neurons_in_the_third_hidden_dense_layer != 0:\n",
    "            list_of_tower_layer_input_and_output_dimensions.append(number_of_neurons_in_the_third_hidden_dense_layer)\n",
    "\n",
    "        return list_of_tower_layer_input_and_output_dimensions\n",
    "\n",
    "    @classmethod\n",
    "    def _create_dense_sequential_stack(cls, list_of_layer_input_and_output_dimensions):\n",
    "        modules = []\n",
    "        number_of_layers = len(list_of_layer_input_and_output_dimensions) - 1\n",
    "        for layer_idx in range(0, number_of_layers):\n",
    "            modules.append(\n",
    "                nn.Linear(\n",
    "                    in_features=list_of_layer_input_and_output_dimensions[layer_idx],\n",
    "                    out_features=list_of_layer_input_and_output_dimensions[layer_idx + 1],\n",
    "                )\n",
    "            )\n",
    "            if layer_idx != (number_of_layers - 1):\n",
    "                modules.append(nn.ReLU())\n",
    "\n",
    "        return nn.Sequential(*modules)\n",
    "\n",
    "    def calculate_item_tower_input_dimension(self):\n",
    "        return (\n",
    "                self._article_id_embedding_dimension\n",
    "                + self._article_garment_group_name_unique_value_count\n",
    "                + self._article_index_name_unique_value_count\n",
    "                + self._article_product_group_name_unique_value_count\n",
    "                + self._article_graphical_appearance_name_unique_value_count\n",
    "                + self._article_perceived_color_master_name_unique_value_count\n",
    "        )\n",
    "\n",
    "    def _create_query_tower_model(\n",
    "            self,\n",
    "            dimension_of_input_for_first_dense_layer,\n",
    "            number_of_neurons_in_the_first_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_second_hidden_dense_layer,\n",
    "            number_of_neurons_in_the_third_hidden_dense_layer,\n",
    "    ):\n",
    "        self.query_tower_model = self._create_dense_sequential_stack(\n",
    "            list_of_layer_input_and_output_dimensions=self._assemble_list_of_sequential_dense_linear_stack_dimensions(\n",
    "                dimension_of_input_for_first_dense_layer=dimension_of_input_for_first_dense_layer,\n",
    "                number_of_neurons_in_the_first_hidden_dense_layer=number_of_neurons_in_the_first_hidden_dense_layer,\n",
    "                number_of_neurons_in_the_second_hidden_dense_layer=number_of_neurons_in_the_second_hidden_dense_layer,\n",
    "                number_of_neurons_in_the_third_hidden_dense_layer=number_of_neurons_in_the_third_hidden_dense_layer,\n",
    "            ),\n",
    "        )\n",
    "\n",
    "    def calculate_query_tower_input_dimension(self):\n",
    "        return (self._customer_id_embedding_dimension + self._customer_transaction_date_dimension\n",
    "                + self._customer_age_group_unique_value_count)\n",
    "\n",
    "    def forward(\n",
    "            self,\n",
    "            article_id_index=None,\n",
    "            article_index_name_index=None,\n",
    "            article_garment_group_name_index=None,\n",
    "            article_graphical_appearance_name_index=None,\n",
    "            article_perceived_color_master_name_index=None,\n",
    "            article_product_group_name_index=None,\n",
    "            customer_id_index=None,\n",
    "            customer_age_group_index=None,\n",
    "            transaction_date_index=None,\n",
    "    ):\n",
    "        item_tower_input = self._preprocess_item_tower_input_using_input_layer(\n",
    "            article_id_index=article_id_index,\n",
    "            article_index_name_index=article_index_name_index,\n",
    "            article_garment_group_name_index=article_garment_group_name_index,\n",
    "            article_product_group_name_index=article_product_group_name_index,\n",
    "            perceived_color_master_name_index=article_perceived_color_master_name_index,\n",
    "            graphical_appearance_name_index=article_graphical_appearance_name_index,\n",
    "        )\n",
    "        query_tower_input = self._preprocess_query_tower_input_using_input_layer(\n",
    "            customer_id_index=customer_id_index,\n",
    "            customer_age_group_index=customer_age_group_index,\n",
    "            transaction_date_index=transaction_date_index,\n",
    "        )\n",
    "\n",
    "        item_embedding = self.item_tower_model(item_tower_input)\n",
    "        query_embedding = self.query_tower_model(query_tower_input)\n",
    "\n",
    "        return self.similarity_function(item_embedding, query_embedding)\n",
    "\n",
    "    def _preprocess_item_tower_input_using_input_layer(\n",
    "            self,\n",
    "            article_id_index,\n",
    "            article_index_name_index,\n",
    "            article_garment_group_name_index,\n",
    "            article_product_group_name_index,\n",
    "            perceived_color_master_name_index,\n",
    "            graphical_appearance_name_index,\n",
    "    ):\n",
    "        feature_tensors = []\n",
    "\n",
    "        if article_id_index is not None and self.article_id_encoder is not None:\n",
    "            feature_tensors.append(self.article_id_encoder(article_id_index.to(self._device)))\n",
    "\n",
    "        if article_index_name_index is not None and self._article_index_name_unique_value_count != 0:\n",
    "            feature_tensors.append(\n",
    "                nn.functional.one_hot(article_index_name_index, self._article_index_name_unique_value_count)\n",
    "            )\n",
    "\n",
    "        if article_garment_group_name_index is not None and self._article_garment_group_name_unique_value_count != 0:\n",
    "            feature_tensors.append(\n",
    "                nn.functional.one_hot(\n",
    "                    article_garment_group_name_index,\n",
    "                    self._article_garment_group_name_unique_value_count,\n",
    "                )\n",
    "            )\n",
    "\n",
    "        if article_product_group_name_index is not None and self._article_product_group_name_unique_value_count != 0:\n",
    "            feature_tensors.append(\n",
    "                nn.functional.one_hot(\n",
    "                    article_product_group_name_index,\n",
    "                    self._article_product_group_name_unique_value_count,\n",
    "                )\n",
    "            )\n",
    "\n",
    "        if perceived_color_master_name_index is not None and self._article_perceived_color_master_name_unique_value_count != 0:\n",
    "            feature_tensors.append(\n",
    "                nn.functional.one_hot(\n",
    "                    perceived_color_master_name_index,\n",
    "                    self._article_perceived_color_master_name_unique_value_count,\n",
    "                )\n",
    "            )\n",
    "\n",
    "        if graphical_appearance_name_index is not None and self._article_graphical_appearance_name_unique_value_count != 0:\n",
    "            feature_tensors.append(\n",
    "                nn.functional.one_hot(\n",
    "                    graphical_appearance_name_index,\n",
    "                    self._article_graphical_appearance_name_unique_value_count,\n",
    "                )\n",
    "            )\n",
    "\n",
    "        return torch.cat(feature_tensors, dim=1)\n",
    "\n",
    "    def _preprocess_query_tower_input_using_input_layer(\n",
    "            self,\n",
    "            customer_id_index,\n",
    "            customer_age_group_index,\n",
    "            transaction_date_index,\n",
    "    ):\n",
    "        feature_tensors = []\n",
    "\n",
    "        if customer_id_index is not None and self.customer_id_encoder is not None:\n",
    "            feature_tensors.append(self.customer_id_encoder(customer_id_index))\n",
    "\n",
    "        if customer_age_group_index is not None and self._customer_age_group_unique_value_count != 0:\n",
    "            feature_tensors.append(\n",
    "                nn.functional.one_hot(customer_age_group_index, self._customer_age_group_unique_value_count))\n",
    "\n",
    "        if transaction_date_index is not None and self.transaction_date_encoder is not None:\n",
    "            feature_tensors.append(self.transaction_date_encoder(transaction_date_index))\n",
    "\n",
    "        return torch.cat(feature_tensors, dim=1)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4f9be7e6ed2c64c0",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The function in the cell below can be used to train the model."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "25d2c8fb3280aa84"
  },
  {
   "cell_type": "code",
   "source": [
    "def train_model(\n",
    "        device,\n",
    "        negative_batch_dataloader,\n",
    "        model,\n",
    "        loss_fn,\n",
    "        optimizer,\n",
    "        article_id_input_key,\n",
    "        article_index_name_input_key,\n",
    "        article_garment_group_input_key,\n",
    "        article_product_group_input_key,\n",
    "        article_perceived_color_master_input_key,\n",
    "        article_graphical_appearance_input_key,\n",
    "        customer_id_input_key,\n",
    "        customer_age_group_input_key,\n",
    "        transaction_date_input_key,\n",
    "):\n",
    "    model.train()\n",
    "    size = len(negative_batch_dataloader.dataset)\n",
    "    for batch_index, (X, y) in enumerate(negative_batch_dataloader):\n",
    "        concatenated_positive_and_negative_items_input = {}\n",
    "        for input_key in X[0].keys():\n",
    "            concatenated_positive_and_negative_items_input[input_key] = torch.cat(\n",
    "                (X[0][input_key], X[1][input_key]),\n",
    "                dim=0,\n",
    "            ).to(device)\n",
    "        concatenated_positive_and_negative_items_label = torch.cat((y[0], y[1]), dim=0).to(device)\n",
    "        logits = model(\n",
    "            article_id_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_id_input_key\n",
    "            ),\n",
    "            article_index_name_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_index_name_input_key\n",
    "            ),\n",
    "            article_garment_group_name_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_garment_group_input_key\n",
    "            ),\n",
    "            article_graphical_appearance_name_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_graphical_appearance_input_key\n",
    "            ),\n",
    "            article_perceived_color_master_name_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_perceived_color_master_input_key\n",
    "            ),\n",
    "            article_product_group_name_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_product_group_input_key\n",
    "            ),\n",
    "            customer_id_index=concatenated_positive_and_negative_items_input.get(\n",
    "                customer_id_input_key\n",
    "            ),\n",
    "            customer_age_group_index=concatenated_positive_and_negative_items_input.get(\n",
    "                customer_age_group_input_key\n",
    "            ),\n",
    "            transaction_date_index=concatenated_positive_and_negative_items_input.get(\n",
    "                transaction_date_input_key\n",
    "            ),\n",
    "        )\n",
    "        loss = loss_fn(logits, concatenated_positive_and_negative_items_label)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        if batch_index % 10 == 0:\n",
    "            loss_value, current = loss.item(), (batch_index + 1) * len(concatenated_positive_and_negative_items_input)\n",
    "            print(f\"Training loss: {loss_value:>7f}  [{current:>5d}/{size:>5d}]\")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "aa161ab3c4f6f25b",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The **perform_loss_test** function can be used to test the model with a loss function. The function calculates average loss."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "3afad323c4d58d8c"
  },
  {
   "cell_type": "code",
   "source": [
    "def perform_loss_test(\n",
    "        device,\n",
    "        negative_batch_dataloader,\n",
    "        model,\n",
    "        loss_fn,\n",
    "        article_id_input_key,\n",
    "        article_index_name_input_key,\n",
    "        article_garment_group_input_key,\n",
    "        article_product_group_input_key,\n",
    "        article_perceived_color_master_input_key,\n",
    "        article_graphical_appearance_input_key,\n",
    "        customer_id_input_key,\n",
    "        customer_age_group_input_key,\n",
    "        transaction_date_input_key,\n",
    "):\n",
    "    loss_test_average_test_loss = 0\n",
    "    number_of_samples_in_dataset = len(negative_batch_dataloader.dataset)\n",
    "    for batch_index, (X, y) in enumerate(negative_batch_dataloader):\n",
    "        concatenated_positive_and_negative_items_input = {}\n",
    "        for input_key in X[0].keys():\n",
    "            concatenated_positive_and_negative_items_input[input_key] = torch.cat(\n",
    "                (X[0][input_key], X[1][input_key]),\n",
    "                dim=0,\n",
    "            ).to(device)\n",
    "        concatenated_positive_and_negative_items_label = torch.cat((y[0], y[1]), dim=0).to(device)\n",
    "        logits = model(\n",
    "            article_id_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_id_input_key\n",
    "            ),\n",
    "            article_index_name_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_index_name_input_key\n",
    "            ),\n",
    "            article_garment_group_name_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_garment_group_input_key\n",
    "            ),\n",
    "            article_graphical_appearance_name_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_graphical_appearance_input_key\n",
    "            ),\n",
    "            article_perceived_color_master_name_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_perceived_color_master_input_key\n",
    "            ),\n",
    "            article_product_group_name_index=concatenated_positive_and_negative_items_input.get(\n",
    "                article_product_group_input_key\n",
    "            ),\n",
    "            customer_id_index=concatenated_positive_and_negative_items_input.get(\n",
    "                customer_id_input_key\n",
    "            ),\n",
    "            customer_age_group_index=concatenated_positive_and_negative_items_input.get(\n",
    "                customer_age_group_input_key\n",
    "            ),\n",
    "            transaction_date_index=concatenated_positive_and_negative_items_input.get(\n",
    "                transaction_date_input_key\n",
    "            ),\n",
    "        )\n",
    "        loss = loss_fn(logits, concatenated_positive_and_negative_items_label)\n",
    "        loss_test_average_test_loss += loss.item() * len(concatenated_positive_and_negative_items_input)\n",
    "\n",
    "        if batch_index % 100 == 0:\n",
    "            current_sample_idx = (batch_index + 1) * len(concatenated_positive_and_negative_items_input)\n",
    "            print(f\"Current average test loss: {loss_test_average_test_loss:>7f}\" +\n",
    "                  f\"[{current_sample_idx:>5d}/{number_of_samples_in_dataset:>5d}]\")\n",
    "\n",
    "    loss_test_average_test_loss /= number_of_samples_in_dataset\n",
    "\n",
    "    return loss_test_average_test_loss"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "dbf0c63946228823",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The **perform_recall_test** tests the model's recall score on a test dataset. The recall scores are calculated for each customer, then calculates the average of these scores and returns it. If a customer has more than k transaction, then the k transactions with the highest score will be used in the calculations."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "3046f5269453e51a"
  },
  {
   "cell_type": "code",
   "source": [
    "def perform_recall_test(\n",
    "        device,\n",
    "        inspected_top_k_per_customer,\n",
    "        positive_interactions_dataloader,\n",
    "        model,\n",
    "        article_id_input_key,\n",
    "        article_index_name_input_key,\n",
    "        article_garment_group_input_key,\n",
    "        article_product_group_input_key,\n",
    "        article_perceived_color_master_input_key,\n",
    "        article_graphical_appearance_input_key,\n",
    "        customer_id_input_key,\n",
    "        customer_age_group_input_key,\n",
    "        transaction_date_input_key,\n",
    "):\n",
    "    highest_logits_of_customers = {}\n",
    "    recall_scores_of_customers = []\n",
    "    number_of_samples_in_dataset = len(positive_interactions_dataloader.dataset)\n",
    "    for batch_idx, (X, y) in enumerate(positive_interactions_dataloader):\n",
    "        logits = model(\n",
    "            article_id_index=X.get(\n",
    "                article_id_input_key\n",
    "            ).to(device),\n",
    "            article_index_name_index=X.get(\n",
    "                article_index_name_input_key\n",
    "            ).to(device),\n",
    "            article_garment_group_name_index=X.get(\n",
    "                article_garment_group_input_key\n",
    "            ).to(device),\n",
    "            article_graphical_appearance_name_index=X.get(\n",
    "                article_graphical_appearance_input_key\n",
    "            ).to(device),\n",
    "            article_perceived_color_master_name_index=X.get(\n",
    "                article_perceived_color_master_input_key\n",
    "            ).to(device),\n",
    "            article_product_group_name_index=X.get(\n",
    "                article_product_group_input_key\n",
    "            ).to(device),\n",
    "            customer_id_index=X.get(\n",
    "                customer_id_input_key\n",
    "            ).to(device),\n",
    "            customer_age_group_index=X.get(\n",
    "                customer_age_group_input_key\n",
    "            ).to(device),\n",
    "            transaction_date_index=X.get(\n",
    "                transaction_date_input_key\n",
    "            ).to(device),\n",
    "        )\n",
    "        for logit_idx, logit in enumerate(logits):\n",
    "            customer_id = X[customer_id_input_key][logit_idx]\n",
    "\n",
    "            if customer_id not in highest_logits_of_customers.keys():\n",
    "                highest_logits_of_customers[customer_id] = []\n",
    "\n",
    "            bisect.insort(highest_logits_of_customers[customer_id], logit, key=lambda x: -x)\n",
    "\n",
    "            if len(highest_logits_of_customers) > inspected_top_k_per_customer:\n",
    "                highest_logits_of_customers[customer_id] \\\n",
    "                    = highest_logits_of_customers[customer_id][:inspected_top_k_per_customer]\n",
    "\n",
    "        if batch_idx % 100 == 0:\n",
    "            current_sample_idx = (batch_idx + 1) * len(X)\n",
    "            print(f\"[{current_sample_idx:>5d}/{number_of_samples_in_dataset:>5d}]\")\n",
    "\n",
    "    for customer_id in highest_logits_of_customers.keys():\n",
    "        true_positive_counter = 0\n",
    "        for logit in highest_logits_of_customers[customer_id]:\n",
    "            if logit >= 0.5:\n",
    "                true_positive_counter += 1\n",
    "        recall_scores_of_customers.append(true_positive_counter / len(highest_logits_of_customers[customer_id]))\n",
    "\n",
    "    return sum(recall_scores_of_customers) / len(recall_scores_of_customers)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "34a921b475fd4913",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The **test_model** functions runs the tests (loss, recall) on the model, and returns their scores."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4897979b439b5cf2"
  },
  {
   "cell_type": "code",
   "source": [
    "def test_model(\n",
    "        device,\n",
    "        loss_test_negative_batch_dataloader,\n",
    "        recall_test_dataloader,\n",
    "        model,\n",
    "        loss_fn,\n",
    "        article_id_input_key,\n",
    "        article_index_name_input_key,\n",
    "        article_garment_group_input_key,\n",
    "        article_product_group_input_key,\n",
    "        article_perceived_color_master_input_key,\n",
    "        article_graphical_appearance_input_key,\n",
    "        customer_id_input_key,\n",
    "        customer_age_group_input_key,\n",
    "        transaction_date_input_key,\n",
    "):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        loss_test_result_average_loss = perform_loss_test(\n",
    "            device=device,\n",
    "            negative_batch_dataloader=loss_test_negative_batch_dataloader,\n",
    "            model=model,\n",
    "            loss_fn=loss_fn,\n",
    "            article_id_input_key=article_id_input_key,\n",
    "            article_index_name_input_key=article_index_name_input_key,\n",
    "            article_garment_group_input_key=article_garment_group_input_key,\n",
    "            article_product_group_input_key=article_product_group_input_key,\n",
    "            article_perceived_color_master_input_key=article_perceived_color_master_input_key,\n",
    "            article_graphical_appearance_input_key=article_graphical_appearance_input_key,\n",
    "            customer_id_input_key=customer_id_input_key,\n",
    "            customer_age_group_input_key=customer_age_group_input_key,\n",
    "            transaction_date_input_key=transaction_date_input_key,\n",
    "        )\n",
    "        recall_score = perform_recall_test(\n",
    "            device=device,\n",
    "            inspected_top_k_per_customer=500,\n",
    "            positive_interactions_dataloader=recall_test_dataloader,\n",
    "            model=model,\n",
    "            article_id_input_key=article_id_input_key,\n",
    "            article_index_name_input_key=article_index_name_input_key,\n",
    "            article_garment_group_input_key=article_garment_group_input_key,\n",
    "            article_product_group_input_key=article_product_group_input_key,\n",
    "            article_perceived_color_master_input_key=article_perceived_color_master_input_key,\n",
    "            article_graphical_appearance_input_key=article_graphical_appearance_input_key,\n",
    "            customer_id_input_key=customer_id_input_key,\n",
    "            customer_age_group_input_key=customer_age_group_input_key,\n",
    "            transaction_date_input_key=transaction_date_input_key,\n",
    "        )\n",
    "\n",
    "        return loss_test_result_average_loss, recall_score"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7c8fce8a2660adc3",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "**load_training_and_test_data** function loads all the data that is required for the model inspection. ~The training dataset contains the data of the first week of 2019, while the test data contains the second week.~ Currently, the training and test dataset contains only 1-1 days from the original dataset; the reason for is that currently I lack the required hardware to run experiments with bigger datasets, but hopefully I can solve this problem in the near future, and then I'll improve this experiment container codebase. Only such articles' and customers' data is kept, that has a transaction in the inspected period. To run this function, the [H&M Personalized Fashion Recommendations dataset](https://www.kaggle.com/competitions/h-and-m-personalized-fashion-recommendations/) needs to be downloaded first. Then change the values of *path_to_articles*, *path_to_customers* and *path_to_transactions* variables to the path of the downloaded dataset."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "850f79dd567d5c41"
  },
  {
   "cell_type": "code",
   "source": [
    "def load_training_and_test_data(\n",
    "        article_id_column_name,\n",
    "        article_garment_group_column_name,\n",
    "        article_graphical_appearance_column_name,\n",
    "        article_index_name_column_name,\n",
    "        article_perceived_color_master_column_name,\n",
    "        article_product_group_column_name,\n",
    "        customer_id_column_name,\n",
    "        customer_age_column_name,\n",
    "        transaction_date_column_name,\n",
    "        customer_age_group_size,\n",
    "        customer_age_minimum,\n",
    "        is_article_garment_group_name_used,\n",
    "        is_article_graphical_appearance_name_used,\n",
    "        is_article_index_name_used,\n",
    "        is_article_perceived_color_master_name_used,\n",
    "        is_article_product_group_name_used,\n",
    "        transaction_date_embedding_dimension,\n",
    "):\n",
    "    path_to_articles = \"C:\\\\Egyetem\\\\MSc\\\\Onallo\\\\HM_dataset\\\\articles.csv\"\n",
    "    path_to_customers = \"C:\\\\Egyetem\\\\MSc\\\\Onallo\\\\HM_dataset\\\\customers.csv\"\n",
    "    path_to_transactions = \"C:\\\\Egyetem\\\\MSc\\\\Onallo\\\\HM_dataset\\\\transactions_train.csv\"\n",
    "\n",
    "    transactions_df_columns = [article_id_column_name, customer_id_column_name, transaction_date_column_name]\n",
    "\n",
    "    transactions_df_iter = pd.read_csv(\n",
    "        filepath_or_buffer=path_to_transactions,\n",
    "        usecols=transactions_df_columns,\n",
    "        iterator=True,\n",
    "        chunksize=100_000,\n",
    "    )\n",
    "\n",
    "    transactions_start_date = \"2019-01-07\"\n",
    "    transactions_end_date = \"2019-01-08\"\n",
    "    test_transactions_delimiter_date = \"2019-01-08\"\n",
    "\n",
    "    transactions_df = pd.concat(\n",
    "        [\n",
    "            transactions_df_chunk[\n",
    "                (transactions_df_chunk[transaction_date_column_name] >= transactions_start_date) &\n",
    "                (transactions_df_chunk[transaction_date_column_name] <= transactions_end_date)\n",
    "                ]\n",
    "            for transactions_df_chunk in transactions_df_iter]\n",
    "    )\n",
    "\n",
    "    # TODO remove head limits after processing is optimized\n",
    "    training_transactions_df = transactions_df[\n",
    "        transactions_df[transaction_date_column_name] < test_transactions_delimiter_date\n",
    "    ].head(1000)\n",
    "    loss_test_transactions_df = transactions_df[\n",
    "        transactions_df[transaction_date_column_name] >= test_transactions_delimiter_date\n",
    "    ].head(1000)\n",
    "    recall_test_transactions_df = loss_test_transactions_df\n",
    "\n",
    "    if transaction_date_embedding_dimension != 0:\n",
    "        transactions_df.drop(transaction_date_column_name, axis=1)\n",
    "        loss_test_transactions_df.drop(transaction_date_column_name, axis=1)\n",
    "        recall_test_transactions_df.drop(transaction_date_column_name, axis=1)\n",
    "\n",
    "    articles_df_columns = [article_id_column_name]\n",
    "\n",
    "    if is_article_garment_group_name_used:\n",
    "        articles_df_columns.append(article_garment_group_column_name)\n",
    "\n",
    "    if is_article_graphical_appearance_name_used:\n",
    "        articles_df_columns.append(article_graphical_appearance_column_name)\n",
    "\n",
    "    if is_article_index_name_used:\n",
    "        articles_df_columns.append(article_index_name_column_name)\n",
    "\n",
    "    if is_article_perceived_color_master_name_used:\n",
    "        articles_df_columns.append(article_perceived_color_master_column_name)\n",
    "\n",
    "    if is_article_product_group_name_used:\n",
    "        articles_df_columns.append(article_product_group_column_name)\n",
    "\n",
    "    articles_df = pd.read_csv(\n",
    "        filepath_or_buffer=path_to_articles,\n",
    "        usecols=articles_df_columns,\n",
    "    )\n",
    "\n",
    "    articles_df = articles_df[articles_df[article_id_column_name].isin(transactions_df[article_id_column_name])]\n",
    "\n",
    "    customers_df_columns = [customer_id_column_name]\n",
    "\n",
    "    if customer_age_group_size != 0:\n",
    "        customers_df_columns.append(customer_age_column_name)\n",
    "\n",
    "    customers_df = pd.read_csv(\n",
    "        filepath_or_buffer=path_to_customers,\n",
    "        usecols=customers_df_columns,\n",
    "    )\n",
    "\n",
    "    customers_df = customers_df[customers_df[customer_id_column_name].isin(transactions_df[customer_id_column_name])]\n",
    "\n",
    "    customers_df[customer_age_column_name].fillna(\n",
    "        value=customers_df[customer_age_column_name].max() + customer_age_minimum / 2,\n",
    "        inplace=True,\n",
    "    )\n",
    "\n",
    "    print(f\"Number of articles: {len(articles_df)}\")\n",
    "    print(f\"Number of customers:  {len(customers_df)}\")\n",
    "    print(f\"Number of training transactions: {len(training_transactions_df)}\")\n",
    "    print(f\"Number of loss test and recall test transactions: {len(loss_test_transactions_df)}\")\n",
    "\n",
    "    return articles_df, customers_df, training_transactions_df, loss_test_transactions_df, recall_test_transactions_df"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "88e8b35736427c02",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "**get_pytorch_device** function determines what kind of hardware will be used for the model's training and testing."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "1c674fe713475ce4"
  },
  {
   "cell_type": "code",
   "source": [
    "def get_pytorch_device():\n",
    "    device = (\n",
    "        \"cuda\"\n",
    "        if torch.cuda.is_available()\n",
    "        else \"mps\"\n",
    "        if torch.backends.mps.is_available()\n",
    "        else \"cpu\"\n",
    "    )\n",
    "\n",
    "    print(f\"Using {device} device\")\n",
    "\n",
    "    return device"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "de4ef25dbed9d118",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "This **train_and_test_model** runs the training and tests of the model in epoch iterations. *NegativeBatchSampledTransactionDataset* doubles the number of items in a single batch because of the negative transaction generation, so only half batch size value is passed to actually use the required batch size."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e45f75457fef0a0"
  },
  {
   "cell_type": "code",
   "source": [
    "def train_and_test_model(\n",
    "        device,\n",
    "        model,\n",
    "        loss_fn,\n",
    "        optimizer,\n",
    "        articles_df,\n",
    "        customers_df,\n",
    "        training_transactions_df,\n",
    "        loss_test_transactions_df,\n",
    "        recall_test_transactions_df,\n",
    "        customer_age_group_size,\n",
    "        customer_minimum_age,\n",
    "        epochs,\n",
    "        training_batch_size,\n",
    "        test_batch_size,\n",
    "        article_id_column_name,\n",
    "        article_garment_group_column_name,\n",
    "        article_graphical_appearance_column_name,\n",
    "        article_index_name_column_name,\n",
    "        article_perceived_color_master_column_name,\n",
    "        article_product_group_column_name,\n",
    "        customer_id_column_name,\n",
    "        customer_age_column_name,\n",
    "        transaction_date_column_name,\n",
    "):\n",
    "    print(\"Starting training and testing...\")\n",
    "    loss_test_loss = 0\n",
    "    recall_score = 0\n",
    "    for t in range(epochs):\n",
    "        print(f\"Epoch {t + 1}\\n-------------------------------\")\n",
    "\n",
    "        train_model(\n",
    "            device=device,\n",
    "            negative_batch_dataloader=DataLoader(\n",
    "                dataset=NegativeBatchSampledTransactionDataset(\n",
    "                    articles_df=articles_df,\n",
    "                    customers_df=customers_df,\n",
    "                    transactions_df=training_transactions_df,\n",
    "                    customer_age_group_size=customer_age_group_size,\n",
    "                    customer_minimum_age=customer_minimum_age,\n",
    "                    article_id_column_name=article_id_column_name,\n",
    "                    article_garment_group_column_name=article_garment_group_column_name,\n",
    "                    article_graphical_appearance_column_name=article_graphical_appearance_column_name,\n",
    "                    article_index_name_column_name=article_index_name_column_name,\n",
    "                    article_perceived_color_master_column_name=article_perceived_color_master_column_name,\n",
    "                    article_product_group_column_name=article_product_group_column_name,\n",
    "                    customer_id_column_name=customer_id_column_name,\n",
    "                    customer_age_column_name=customer_age_column_name,\n",
    "                    transaction_date_column_name=transaction_date_column_name,\n",
    "                ),\n",
    "                batch_size=round(training_batch_size / 2),\n",
    "            ),\n",
    "            model=model,\n",
    "            loss_fn=loss_fn,\n",
    "            optimizer=optimizer,\n",
    "            article_id_input_key=article_id_column_name,\n",
    "            article_index_name_input_key=article_index_name_column_name,\n",
    "            article_garment_group_input_key=article_garment_group_column_name,\n",
    "            article_product_group_input_key=article_product_group_column_name,\n",
    "            article_perceived_color_master_input_key=article_perceived_color_master_column_name,\n",
    "            article_graphical_appearance_input_key=article_graphical_appearance_column_name,\n",
    "            customer_id_input_key=customer_id_column_name,\n",
    "            customer_age_group_input_key=customer_age_column_name,\n",
    "            transaction_date_input_key=transaction_date_column_name,\n",
    "        )\n",
    "\n",
    "        loss_test_loss, recall_score = test_model(\n",
    "            device=device,\n",
    "            loss_test_negative_batch_dataloader=DataLoader(\n",
    "                dataset=NegativeBatchSampledTransactionDataset(\n",
    "                    articles_df=articles_df,\n",
    "                    customers_df=customers_df,\n",
    "                    transactions_df=loss_test_transactions_df,\n",
    "                    customer_age_group_size=customer_age_group_size,\n",
    "                    customer_minimum_age=customer_minimum_age,\n",
    "                    article_id_column_name=article_id_column_name,\n",
    "                    article_garment_group_column_name=article_garment_group_column_name,\n",
    "                    article_graphical_appearance_column_name=article_graphical_appearance_column_name,\n",
    "                    article_index_name_column_name=article_index_name_column_name,\n",
    "                    article_perceived_color_master_column_name=article_perceived_color_master_column_name,\n",
    "                    article_product_group_column_name=article_product_group_column_name,\n",
    "                    customer_id_column_name=customer_id_column_name,\n",
    "                    customer_age_column_name=customer_age_column_name,\n",
    "                    transaction_date_column_name=transaction_date_column_name,\n",
    "                ),\n",
    "                batch_size=round(test_batch_size / 2),\n",
    "            ),\n",
    "            recall_test_dataloader=DataLoader(\n",
    "                dataset=TransactionDataset(\n",
    "                    articles_df=articles_df,\n",
    "                    customers_df=customers_df,\n",
    "                    transactions_df=recall_test_transactions_df,\n",
    "                    customer_age_group_size=customer_age_group_size,\n",
    "                    customer_minimum_age=customer_minimum_age,\n",
    "                    article_id_column_name=article_id_column_name,\n",
    "                    article_garment_group_column_name=article_garment_group_column_name,\n",
    "                    article_graphical_appearance_column_name=article_graphical_appearance_column_name,\n",
    "                    article_index_name_column_name=article_index_name_column_name,\n",
    "                    article_perceived_color_master_column_name=article_perceived_color_master_column_name,\n",
    "                    article_product_group_column_name=article_product_group_column_name,\n",
    "                    customer_id_column_name=customer_id_column_name,\n",
    "                    customer_age_column_name=customer_age_column_name,\n",
    "                    transaction_date_column_name=transaction_date_column_name,\n",
    "                ),\n",
    "                batch_size=test_batch_size,\n",
    "            ),\n",
    "            model=model,\n",
    "            loss_fn=loss_fn,\n",
    "            article_id_input_key=article_id_column_name,\n",
    "            article_index_name_input_key=article_index_name_column_name,\n",
    "            article_garment_group_input_key=article_garment_group_column_name,\n",
    "            article_product_group_input_key=article_product_group_column_name,\n",
    "            article_perceived_color_master_input_key=article_perceived_color_master_column_name,\n",
    "            article_graphical_appearance_input_key=article_graphical_appearance_column_name,\n",
    "            customer_id_input_key=customer_id_column_name,\n",
    "            customer_age_group_input_key=customer_age_column_name,\n",
    "            transaction_date_input_key=transaction_date_column_name,\n",
    "        )\n",
    "\n",
    "    print(f\"Loss test loss: {loss_test_loss}\")\n",
    "    print(f\"Recall test recall: {recall_score}\")\n",
    "    return loss_test_loss, recall_score"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "db17ef6544b38e3",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "**perform_configured_model_inspection** function uses its *config* parameter to initialize the model inspections parameters, and then runs it."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4cae8baa39d78719"
  },
  {
   "cell_type": "code",
   "source": [
    "def perform_configured_model_inspection(config, device):\n",
    "    article_id_column_name = \"article_id\"\n",
    "    article_garment_group_column_name = \"garment_group_name\"\n",
    "    article_graphical_appearance_column_name = \"graphical_appearance_name\"\n",
    "    article_index_name_column_name = \"index_name\"\n",
    "    article_perceived_color_master_column_name = \"perceived_colour_master_name\"\n",
    "    article_product_group_column_name = \"product_group_name\"\n",
    "    customer_id_column_name = \"customer_id\"\n",
    "    customer_age_column_name = \"age\"\n",
    "    transaction_date_column_name = \"t_dat\"\n",
    "\n",
    "    customer_maximum_age = 99\n",
    "    customer_minimum_age = 16\n",
    "\n",
    "    articles_df, customers_df, training_transactions_df, loss_test_transactions_df, recall_test_transactions_df \\\n",
    "        = load_training_and_test_data(\n",
    "        article_id_column_name=article_id_column_name,\n",
    "        article_garment_group_column_name=article_garment_group_column_name,\n",
    "        article_graphical_appearance_column_name=article_graphical_appearance_column_name,\n",
    "        article_index_name_column_name=article_index_name_column_name,\n",
    "        article_perceived_color_master_column_name=article_perceived_color_master_column_name,\n",
    "        article_product_group_column_name=article_product_group_column_name,\n",
    "        customer_id_column_name=customer_id_column_name,\n",
    "        customer_age_column_name=customer_age_column_name,\n",
    "        transaction_date_column_name=transaction_date_column_name,\n",
    "        customer_age_group_size=config.customer_age_group_size,\n",
    "        customer_age_minimum=customer_minimum_age,\n",
    "        is_article_garment_group_name_used=config.is_article_garment_group_name_used,\n",
    "        is_article_graphical_appearance_name_used=config.is_article_graphical_appearance_name_used,\n",
    "        is_article_index_name_used=config.is_article_index_name_used,\n",
    "        is_article_perceived_color_master_name_used=config.is_article_perceived_color_master_name_used,\n",
    "        is_article_product_group_name_used=config.is_article_product_group_name_used,\n",
    "        transaction_date_embedding_dimension=config.transaction_date_embedding_dimension,\n",
    "    )\n",
    "\n",
    "    two_tower_model = TwoTowerModel(\n",
    "        device=device,\n",
    "        number_of_neurons_in_the_first_hidden_dense_layer=config.number_of_neurons_in_the_first_hidden_dense_layer,\n",
    "        number_of_neurons_in_the_second_hidden_dense_layer=config.number_of_neurons_in_the_second_hidden_dense_layer,\n",
    "        number_of_neurons_in_the_third_hidden_dense_layer=config.number_of_neurons_in_the_third_hidden_dense_layer,\n",
    "        article_id_embedding_dimension_and_unique_value_count=(\n",
    "            config.article_id_embedding_dimension,\n",
    "            len(articles_df)\n",
    "        ) if config.article_id_embedding_dimension != 0 else None,\n",
    "        article_index_name_unique_value_count=(\n",
    "            len(WebshopDataContainer.get_article_index_name_numeric_encoder()) if config.is_article_index_name_used else 0\n",
    "        ),\n",
    "        article_garment_group_name_unique_value_count=(\n",
    "            len(WebshopDataContainer.get_article_garment_group_name_numeric_encoder())\n",
    "            if config.is_article_garment_group_name_used else 0\n",
    "        ),\n",
    "        article_graphical_appearance_name_unique_value_count=(\n",
    "            len(WebshopDataContainer.get_article_graphical_appearance_numeric_encoder())\n",
    "            if config.is_article_graphical_appearance_name_used else 0\n",
    "        ),\n",
    "        article_product_group_name_unique_value_count=(\n",
    "            len(WebshopDataContainer.get_article_product_group_name_numeric_encoder())\n",
    "            if config.is_article_product_group_name_used else 0\n",
    "        ),\n",
    "        article_perceived_color_master_name_unique_value_count=(\n",
    "            len(WebshopDataContainer.get_article_perceived_color_master_name_numeric_encoder())\n",
    "            if config.is_article_perceived_color_master_name_used else 0\n",
    "        ),\n",
    "        customer_id_embedding_dimension_and_unique_value_count=(\n",
    "            config.customer_id_embedding_dimension,\n",
    "            len(customers_df),\n",
    "        ) if config.customer_id_embedding_dimension != 0 else None,\n",
    "        customer_age_group_count=math.ceil(\n",
    "            (customer_maximum_age - customer_minimum_age) / config.customer_age_group_size\n",
    "        ) if config.customer_age_group_size != 0 else 0,\n",
    "        transaction_date_embedding_dimension_and_unique_value_count=(\n",
    "            config.transaction_date_embedding_dimension,\n",
    "            training_transactions_df[\"t_dat\"].nunique(),\n",
    "        ) if config.transaction_date_embedding_dimension != 0 else None,\n",
    "        similarity_function=nn.CosineSimilarity(),\n",
    "    ).to(device)\n",
    "    cross_entropy_loss_fn = nn.CrossEntropyLoss()\n",
    "    two_tower_model_adam_optimizer = torch.optim.Adam(\n",
    "        params=two_tower_model.parameters(),\n",
    "        lr=config.learning_rate,\n",
    "        betas=(config.beta1, config.beta2),\n",
    "        weight_decay=config.weight_decay,\n",
    "        amsgrad=config.is_amsgrad_used,\n",
    "    )\n",
    "\n",
    "    test_loss, recall_score = train_and_test_model(\n",
    "        device = device,\n",
    "        model=two_tower_model,\n",
    "        loss_fn=cross_entropy_loss_fn,\n",
    "        optimizer=two_tower_model_adam_optimizer,\n",
    "        articles_df=articles_df,\n",
    "        customers_df=customers_df,\n",
    "        training_transactions_df=training_transactions_df,\n",
    "        loss_test_transactions_df=loss_test_transactions_df,\n",
    "        recall_test_transactions_df=recall_test_transactions_df,\n",
    "        customer_age_group_size=config.customer_age_group_size,\n",
    "        customer_minimum_age=customer_minimum_age,\n",
    "        epochs=config.epochs,\n",
    "        training_batch_size=config.batch_size,\n",
    "        test_batch_size=config.batch_size,\n",
    "        article_id_column_name=article_id_column_name,\n",
    "        article_garment_group_column_name=article_garment_group_column_name,\n",
    "        article_graphical_appearance_column_name=article_graphical_appearance_column_name,\n",
    "        article_index_name_column_name=article_index_name_column_name,\n",
    "        article_perceived_color_master_column_name=article_perceived_color_master_column_name,\n",
    "        article_product_group_column_name=article_product_group_column_name,\n",
    "        customer_id_column_name=customer_id_column_name,\n",
    "        customer_age_column_name=customer_age_column_name,\n",
    "        transaction_date_column_name=transaction_date_column_name,\n",
    "    )\n",
    "\n",
    "    return two_tower_model, test_loss, recall_score"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "62a49d53751718a",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "**perform_configured_model_inspection_with_exception_logging** wraps *perform_configured_model_inspection* into a try-catch block, and logs the exception's trace. The main reason of this function's existence to detect errors in the model inspection's implementation."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "32e056ca6f8f3df5"
  },
  {
   "cell_type": "code",
   "source": [
    "def perform_configured_model_inspection_with_exception_logging(config, device):\n",
    "    try:\n",
    "        return perform_configured_model_inspection(config=config, device=device)\n",
    "    except Exception as exception:\n",
    "        print(\"Model inspection was interrupted due to an exception being thrown:\")\n",
    "        traceback.print_exc()\n",
    "        raise exception"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2dcdd1064392262c",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "**save_model_in_onnx_format_and_get_path** saves the received PyTorch model to ONNX format, and returns the saved model's local path. The function uses the relatively old *torch.onnx.export* function instead the newer *torch.onnx.dynamo_export* function, because unfortunately the newer solution didn't work when I tried to use it. The ONNX model's input shape saved with dynamic dimension value, because it allows to use the model with parameterizable batch size."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c8ee7d737142f074"
  },
  {
   "cell_type": "code",
   "source": [
    "def save_model_in_onnx_format_and_get_path(model, filename, input_dimension, sample_input=None):\n",
    "    export_device = \"cpu\"\n",
    "    \n",
    "    local_destination_path = f\"../ml_artifacts/{filename}.onnx\"\n",
    "    if sample_input is None:\n",
    "        sample_input = torch.randn(1, input_dimension).to(export_device)\n",
    "    torch.onnx.export(\n",
    "        model.to(export_device),\n",
    "        sample_input,\n",
    "        local_destination_path,\n",
    "        export_params=True,\n",
    "        input_names=[\"input\"],\n",
    "        output_names=[\"output\"],\n",
    "        dynamic_axes={\n",
    "            \"input\": {0: \"batch_size\"},\n",
    "            \"output\": {0: \"batch_size\"}\n",
    "        },\n",
    "    )\n",
    "    return local_destination_path"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "25abb712d8d19c3b",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "**save_model_as_wandb_artifact** saves the model as ONNX model, then uploads to wandb."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "97279283395e6fd6"
  },
  {
   "cell_type": "code",
   "source": [
    "def save_model_as_wandb_artifact(model, artifact_model_name, filename, input_dimension, sample_input=None):\n",
    "    model_artifact_path = save_model_in_onnx_format_and_get_path(\n",
    "        model=model,\n",
    "        filename=filename,\n",
    "        input_dimension=input_dimension,\n",
    "        sample_input=sample_input,\n",
    "    )\n",
    "    model_artifact = wandb.Artifact(name=artifact_model_name, type=\"model\")\n",
    "    model_artifact.add_file(model_artifact_path)\n",
    "    wandb.log_artifact(model_artifact)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7052ea318e794f1f",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The model inspections and hyperparameter optimizations are run using [wand sweep](https://docs.wandb.ai/guides/sweeps). This function runs the model examinations, then uploads the score of the inspected model and saves and uploads the tower models and embedding layers to wandb."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e9052e381177d11f"
  },
  {
   "cell_type": "code",
   "source": [
    "def sweep_main():\n",
    "    wandb.init()\n",
    "    \n",
    "    device = get_pytorch_device()\n",
    "    model, loss, recall_score = perform_configured_model_inspection_with_exception_logging(device=device, config=wandb.config)\n",
    "    \n",
    "    wandb.log({\n",
    "        \"loss\": loss,\n",
    "        \"recall\": recall_score,\n",
    "    })\n",
    "    save_model_as_wandb_artifact(\n",
    "        model=model.item_tower_model,\n",
    "        artifact_model_name=\"Retrival_two_tower_model_item_tower\",\n",
    "        filename=\"retrieval_item_tower_model\",\n",
    "        input_dimension=model.calculate_item_tower_input_dimension(),\n",
    "    )\n",
    "    save_model_as_wandb_artifact(\n",
    "        model=model.query_tower_model,\n",
    "        artifact_model_name=\"Retrival_two_tower_model_query_tower\",\n",
    "        filename=\"retrieval_query_tower_model\",\n",
    "        input_dimension=model.calculate_query_tower_input_dimension(),\n",
    "    )\n",
    "    save_model_as_wandb_artifact(\n",
    "        model=model.article_id_encoder,\n",
    "        artifact_model_name=\"Retrival_two_tower_model_item_id_encoder\",\n",
    "        filename=\"retrieval_item_id_encoder\",\n",
    "        input_dimension=1,\n",
    "        sample_input=torch.zeros((1, 1), dtype=torch.int64),\n",
    "    )\n",
    "    save_model_as_wandb_artifact(\n",
    "        model=model.customer_id_encoder,\n",
    "        artifact_model_name=\"Retrival_two_tower_model_query_id_encoder\",\n",
    "        filename=\"retrieval_query_id_encoder\",\n",
    "        input_dimension=1,\n",
    "        sample_input=torch.zeros((1, 1), dtype=torch.int64),\n",
    "    )\n",
    "    if model.transaction_date_encoder is not None:\n",
    "        save_model_as_wandb_artifact(\n",
    "            model=model.transaction_date_encoder,\n",
    "            artifact_model_name=\"Retrival_two_tower_query_transaction_date_encoder\",\n",
    "            filename=\"retrieval_query_transaction_date_encoder\",\n",
    "            input_dimension=1,\n",
    "            sample_input=torch.zeros((1, 1), dtype=torch.int64),\n",
    "        )"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7c53bfb71a7bcc75",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "This cell defines the hyperparameter attributes' value range. The syntax is defined by the requirements of wandb sweep. For now except learning rate and optimizer parameters, fixed values are used; soon I'll improve this notebook, and then I will add and support experiments with different parameters."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "3f8155a86adc679b"
  },
  {
   "cell_type": "code",
   "source": [
    "sweep_parameters = {\n",
    "    \"batch_size\": {\"values\": [256], },  # [128, 256, 512, 1024, 2048]\n",
    "    \"epochs\": {\"values\": [1], },  # [1, 2, 4]\n",
    "\n",
    "    \"learning_rate\": {\"max\": 0.1, \"min\": 0.0001, },\n",
    "    \"beta1\": {\"max\": 0.99, \"min\": 0.9, },\n",
    "    \"beta2\": {\"max\": 0.9999, \"min\": 0.999, },\n",
    "    \"weight_decay\": {\"max\": 0.01, \"min\": 0.0001, },\n",
    "    \"is_amsgrad_used\": {\"values\": [0, 1], },\n",
    "\n",
    "    \"number_of_neurons_in_the_first_hidden_dense_layer\": {\"values\": [128], },  # [128, 256]\n",
    "    \"number_of_neurons_in_the_second_hidden_dense_layer\": {\"values\": [64], },  # [64, 128, 256]\n",
    "    \"number_of_neurons_in_the_third_hidden_dense_layer\": {\"values\": [32], },  # [0, 32, 64, 128, 256]\n",
    "\n",
    "    \"article_id_embedding_dimension\": {\"values\": [32], },  # [16, 32, 64, 128]\n",
    "    \"is_article_garment_group_name_used\": {\"values\": [1], },  # [0, 1]\n",
    "    \"is_article_graphical_appearance_name_used\": {\"values\": [1], },  # [0, 1]\n",
    "    \"is_article_index_name_used\": {\"values\": [1], },  # [0, 1]\n",
    "    \"is_article_perceived_color_master_name_used\": {\"values\": [1], },  # [0, 1]\n",
    "    \"is_article_product_group_name_used\": {\"values\": [1], },  # [0, 1]\n",
    "\n",
    "    \"customer_id_embedding_dimension\": {\"values\": [32], },  # [16, 32, 64, 128]\n",
    "    \"customer_age_group_size\": {\"values\": [4], },  # [0, 4, 8]\n",
    "    \"transaction_date_embedding_dimension\": {\"values\": [0], },  # [0, 8, 16, 32, 64, 128, 256]\n",
    "}\n",
    "\n",
    "sweep_configuration = {\n",
    "    \"method\": \"bayes\",\n",
    "    \"name\": \"retrieval-model-sweep\",\n",
    "    \"metric\": {\n",
    "        \"goal\": \"maximize\",\n",
    "        \"name\": \"recall\",\n",
    "    },\n",
    "    \"parameters\": sweep_parameters,\n",
    "}"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6dedfc452247d30f",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The code below authenticates in wandb."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "fd4449409d3592e9"
  },
  {
   "cell_type": "code",
   "source": [
    "wandb.login()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "dd2bc093432e5509",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "This code cell creates a sweep - a parametrized model inspection - and then executes it."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "569baa1a6e86bf75"
  },
  {
   "cell_type": "code",
   "source": [
    "sweep_id = wandb.sweep(sweep=sweep_configuration, project=\"kotlin-clothing-webshop\")\n",
    "\n",
    "wandb.agent(sweep_id=sweep_id, function=sweep_main, count=1)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c55bb815810458da",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 2. Planned further development tasks"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e527f13da5829ee8"
  },
  {
   "cell_type": "markdown",
   "source": [
    "The current inspections uses a relatively small dataset because of the limited hardware that is currently available to me to train and test the model on. I'm actively searching for opportunities to get better hardware because of this reason. After I get better hardware, it could be a good idea to modify the possible value range of some hyperparameters."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7424beb9c6da2e0f"
  },
  {
   "cell_type": "markdown",
   "source": [
    "The current solution could be improved using a [multi-stage recommendation system architecture](https://resources.nvidia.com/en-us-merlin/bad-a-multi-stage-recommender). The currently implemented two tower model could be used as its retrieval phase, and a ranker model could be added to produce better predictions using the output of the retrieval model. Implementing a filtering phase could also make sense, it seems possible that the customers usually don't buy articles bought previously, but first this idea should be checked and confirmed in the [Exploratory Data Analysis Notebook](hm_dataset_inspection_eda.ipynb)."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "a26983ba35f35082"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}