filter_rewards.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "pro_dir = \"./prepared_whole_data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(pro_dir, 'seq_rewards.pkl'), \"rb\") as f:\n",
    "    data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_rewards(seq_trs, rew_trs, seq_tes=None, rew_tes=None):\n",
    "    new_seq_trs, new_seq_tes = [], []\n",
    "    for i in range(len(seq_trs)):\n",
    "        seq_tr = seq_trs[i]\n",
    "        rew_tr = rew_trs[i]\n",
    "        new_seq_tr = []\n",
    "        for j in range(len(seq_tr)):\n",
    "            if rew_tr[j] > 0:\n",
    "                new_seq_tr.append(seq_tr[j])\n",
    "        if not new_seq_tr or len(new_seq_tr) < 2:\n",
    "            continue\n",
    "        if seq_tes and rew_tes:\n",
    "            seq_te = seq_tes[i]\n",
    "            rew_te = rew_tes[i]\n",
    "            new_seq_te = []\n",
    "            for j in range(len(seq_te)):\n",
    "                if rew_te[j] > 0:\n",
    "                    new_seq_te.append(seq_te[j])\n",
    "            if not new_seq_te:\n",
    "                continue\n",
    "            new_seq_tes.append(new_seq_te)\n",
    "        new_seq_trs.append(new_seq_tr)\n",
    "    if new_seq_tes:\n",
    "        return new_seq_trs, new_seq_tes\n",
    "    return new_seq_trs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_data_tr, val_data_te = filter_rewards(data[\"vad_data_tr_seq\"], data[\"vad_data_tr_rewards\"], data[\"vad_data_te_seq\"], data[\"vad_data_te_rewards\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(val_data_tr) == len(val_data_te)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data_tr, test_data_te = filter_rewards(data[\"test_data_tr_seq\"], data[\"test_data_tr_rewards\"], data[\"test_data_te_seq\"], data[\"test_data_te_rewards\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "39456"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(test_data_tr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = filter_rewards(data[\"train_data_seq\"], data[\"train_data_rewards\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "383271"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_to_save = {\n",
    "    \"train\": train_data,\n",
    "    \"validation_tr\": val_data_tr,\n",
    "    \"validation_te\": val_data_te,\n",
    "    \"test_tr\": test_data_tr, \n",
    "    \"test_te\": test_data_te,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(pro_dir, 'seq_dataset.pkl'), \"wb\") as f:\n",
    "    pickle.dump(data_to_save, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "6ea7e837897e166b43df5ffdf98ade3bf2b18becb4ca45509eb8f04aa783dc86"
  },
  "kernelspec": {
   "display_name": "Python 3.8.13 ('venv': venv)",
   "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.8.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}