src/model/generate_embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85155211",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install duckdb==0.7.1 \\\n",
    "duckdb-engine \\\n",
    "watermark \\\n",
    "jupysql \\\n",
    "sqlalchemy \\\n",
    "python-snappy \\\n",
    "pyarrow \\\n",
    "memray \\\n",
    "pandas  \\\n",
    "ipywidgets  \\\n",
    "matplotlib \\\n",
    "gensim \\\n",
    "nltk \\\n",
    "plotly \\\n",
    "redis==4.5.3 \\\n",
    "jupyter-black \\\n",
    "sentence_transformers \\\n",
    "redis \\\n",
    "jupyter_black"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import sys\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import pyarrow as pa\n",
    "import pandas as pd"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a27ff70b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Autoformat cells on run\n",
    "import jupyter_black\n",
    "import pandas as pd\n",
    "\n",
    "jupyter_black.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfb46e93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set log level for model training\n",
    "import logging\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "logging.basicConfig(\n",
    "    format=\"%(levelname)s - %(asctime)s: %(message)s\",\n",
    "    datefmt=\"%H:%M:%S\",\n",
    "    level=logging.INFO,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3e44044",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = pd.read_parquet(\"~/ssl/2023071112_training.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "343fc1c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93c1473e",
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus = embeddings[\"sentence\"].tolist()\n",
    "titles = embeddings[\"title\"].tolist()\n",
    "author = embeddings[\"author\"].tolist()\n",
    "link = embeddings[\"link\"].tolist()\n",
    "review_count = embeddings[\"text_reviews_count\"].tolist()\n",
    "indices = embeddings.index.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d7f04f9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer, util\n",
    "\n",
    "#  A common value for BERT & Co. are 512 word pieces, which correspond to about 300-400 words (for English).\n",
    "# Longer texts than this are truncated to the first x word pieces.\n",
    "# By default, the provided methods use a limit fo 128 word pieces, longer inputs will be truncated\n",
    "# the runtime and the memory requirement grows quadratic with the input length - we'll have to play around with this\n",
    "\n",
    "# Change the length to 200\n",
    "model = SentenceTransformer(\"sentence-transformers/msmarco-distilbert-base-v3\")\n",
    "model.max_seq_length = 200\n",
    "\n",
    "\n",
    "corpus_embeddings = model.encode(\n",
    "    corpus, show_progress_bar=True, device=\"cuda\", convert_to_numpy=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2fa56e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings_list = [x.tolist() for x in corpus_embeddings]\n",
    "embedding_tuple = list(\n",
    "    zip(titles, indices, author, link, review_count, embeddings_list)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9a12ab7",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(\n",
    "    embedding_tuple,\n",
    "    columns=[\"title\", \"index\", \"author\", \"link\", \"review_count\", \"embeddings\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f46d6122",
   "metadata": {},
   "outputs": [],
   "source": [
    "# introspect\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83ff21b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fields = [\n",
    "    (\"title\", pa.string()),\n",
    "    (\"index\", pa.int64(), False),\n",
    "    (\"author\", pa.string()),\n",
    "    (\"link\", pa.string()),\n",
    "    (\"review_count\", pa.string()),\n",
    "    (\"embeddings\", pa.large_list(pa.float64())),\n",
    "]\n",
    "schema = pa.schema(fields)\n",
    "\n",
    "df.to_parquet(\n",
    "    \"20230711_learned_embeddings.snappy\",\n",
    "    engine=\"pyarrow\",\n",
    "    compression=\"snappy\",\n",
    "    schema=schema,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_pytorch_p310)",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}