notebooks/sentencebert-finetuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "470f1f65-5712-42a8-bcf0-2a560ee2baae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import sys\n",
    "sys.path.append('../')\n",
    "import math, statistics, time\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "from datetime import datetime\n",
    "import pickle\n",
    "import pandas as pd\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from sentence_transformers import SentenceTransformer, LoggingHandler, losses, InputExample\n",
    "from tqdm import tqdm\n",
    "from utils.sentence_bert_dataloader import SentenceBertDataloader\n",
    "from utils.dataset import Dataset\n",
    "\n",
    "base_model = 'roberta-base'\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "num_epochs = 20\n",
    "model_name = 'sentence_transformer_30'\n",
    "model_save_path = '../models/{}'.format(model_name)\n",
    "\n",
    "with open('../data/training_label_100.pkl', 'rb') as f:\n",
    "    labels = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a70fc2b4-f09d-4837-bc84-06e537f15bb2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Keys in meme dict dataset: dict_keys(['label_uuid_dic', 'uuid_label_dic', 'uuid_caption_dic', 'uuid_image_path_dic', 'uuid_caption_cased_dic'])\n",
      "Number of uuids: 300\n"
     ]
    }
   ],
   "source": [
    "# load meme dataset\n",
    "meme_dict = None\n",
    "with open('../data/meme_900k_cleaned_data_v2.pkl', 'rb') as f:\n",
    "    meme_dict = pickle.load(f)\n",
    "print(\"Keys in meme dict dataset:\", meme_dict.keys())\n",
    "print(\"Number of uuids:\", len(meme_dict['uuid_label_dic']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "091900ea-261c-4251-9f5b-f31442ac49d8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3cf16c42-aad7-4f3c-8bcd-f21947364fb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# utility functions\n",
    "def clean_and_unify_caption(caption):\n",
    "    return caption[0].strip()+'; '+caption[1].strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2f7c563e-bdf0-4123-9b45-66697b758f25",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "270000 30000\n"
     ]
    }
   ],
   "source": [
    "# create pandas dataframe\n",
    "training_uuids = labels.keys()\n",
    "temp_arr = []\n",
    "for uuid in training_uuids:\n",
    "    for caption in meme_dict['uuid_caption_dic'][uuid]:\n",
    "        temp_arr.append([uuid, clean_and_unify_caption(caption)])\n",
    "df = pd.DataFrame(temp_arr, columns=['category', 'text'])\n",
    "\n",
    "# split dataset\n",
    "np.random.seed(42)\n",
    "df_train, df_test = np.split(df.sample(frac=1, random_state=42), [int(.9*len(df))])\n",
    "\n",
    "print(len(df_train), len(df_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "321c6a26-bab0-4377-a727-54352700fd58",
   "metadata": {},
   "source": [
    "## Creating Dataset and DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7dd9eab1-4fc8-4b01-abb5-398810cb5951",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = Dataset(df_train, labels)\n",
    "test_dataset = Dataset(df_test, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5e1c16e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "202500it [00:00, 733497.27it/s]\n",
      "22500it [00:00, 715014.02it/s]\n"
     ]
    }
   ],
   "source": [
    "train_loader = SentenceBertDataloader(train_dataset, 32)\n",
    "test_loader = SentenceBertDataloader(test_dataset, 32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "208d701f-3e2f-4041-8a83-55648f09a40e",
   "metadata": {},
   "source": [
    "## Model Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "493fe83a-d70e-4c86-8beb-8543aac0d569",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SentenceTransformer('../models/sentence_transformer_roberta_20', device=device)\n",
    "train_loss = losses.ContrastiveLoss(model=model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc85f4af",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(train_objectives=[(train_loader, train_loss)],\n",
    "                              epochs=num_epochs, \n",
    "                              warmup_steps=100, \n",
    "                              output_path=model_save_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1d6f748-f179-45f4-8ce8-c37ed905958f",
   "metadata": {},
   "source": [
    "## Create and save category embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e15c55a5-2f41-4808-892b-cbf8e21373a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def getCategoryEmbeddings(df_train, model):\n",
    "    uuid_to_emb_dict = {}\n",
    "    uuid_count_dict = defaultdict(int)\n",
    "    batch_size = 512\n",
    "    \n",
    "    for i in tqdm(range(0, df_train.shape[0], batch_size)):\n",
    "        texts = list(df_train.text[i:i+batch_size])\n",
    "        uuids = list(df_train.category[i:i+batch_size])\n",
    "        embeddings = model.encode(texts)\n",
    "        for i, uuid in enumerate(uuids):\n",
    "            uuid_count_dict[uuid]+=1\n",
    "            if uuid in uuid_to_emb_dict:\n",
    "                uuid_to_emb_dict[uuid]=uuid_to_emb_dict[uuid]+embeddings[i]\n",
    "            else:\n",
    "                uuid_to_emb_dict[uuid]=embeddings[i]\n",
    "    \n",
    "    for k, v in uuid_to_emb_dict.items():\n",
    "        uuid_to_emb_dict[k] = uuid_to_emb_dict[k]/uuid_count_dict[k] \n",
    "    \n",
    "    return uuid_to_emb_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "31cdd9ee-98a2-4fcf-8eb5-e5c3bf918c2f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 528/528 [02:49<00:00,  3.11it/s]\n",
      "100%|██████████| 528/528 [02:47<00:00,  3.16it/s]\n",
      "100%|██████████| 528/528 [02:48<00:00,  3.14it/s]\n",
      "100%|██████████| 528/528 [02:47<00:00,  3.16it/s]\n",
      "100%|██████████| 528/528 [02:46<00:00,  3.16it/s]\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 5\n",
    "import os\n",
    "for i in range(1,6):\n",
    "    model_name = 'sentence_transformer_roberta_samples_100_epochs_{}'.format(i*5)\n",
    "    model_load_path = '../models/{}'.format(model_name)\n",
    "    embeddings_save_path = '../models/model_utils/{}/category_embeddings.pkl'.format(model_name)\n",
    "    \n",
    "    model = SentenceTransformer(model_load_path, device=device)\n",
    "    uuid_to_emb_dict = getCategoryEmbeddings(df_train, model)\n",
    "    \n",
    "    os.makedirs(os.path.dirname(embeddings_save_path), exist_ok=True)\n",
    "    with open(embeddings_save_path, 'wb') as f:\n",
    "        pickle.dump(uuid_to_emb_dict, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bec4a8a1-00b8-477e-9e10-ec3b62b38e26",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "meme",
   "language": "python",
   "name": "meme"
  },
  "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}