06a_chunk_and_embed.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chunk and embed data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Imports**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Pandarallel will run on 8 workers.\n",
      "INFO: Pandarallel will use standard multiprocessing data transfer (pipe) to transfer data between the main process and workers.\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from pandarallel import pandarallel\n",
    "import os\n",
    "import re\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "pd.options.mode.chained_assignment = None\n",
    "pd.options.display.max_rows = 500\n",
    "pd.options.display.max_seq_items = 500\n",
    "pandarallel.initialize(progress_bar=True)\n",
    "\n",
    "# Suppress Hugginface warning about tokenizers.\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "from staatsarchiv_utils import chunk_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "\n",
    "PREP_OUTPUT_KRP = os.getenv(\"PREP_OUTPUT_KRP\")\n",
    "PREP_OUTPUT_RRB = os.getenv(\"PREP_OUTPUT_RRB\")\n",
    "PREP_OUTPUT_OS = os.getenv(\"PREP_OUTPUT_OS\")\n",
    "PREP_OUTPUT_ABl = os.getenv(\"PREP_OUTPUT_ABl\")\n",
    "\n",
    "DATA_OUTPUT_FULL = os.getenv(\"DATA_OUTPUT_FULL\")\n",
    "DATA_OUTPUT_CHUNKS = os.getenv(\"DATA_OUTPUT_CHUNKS\")\n",
    "DATA_EMBEDDINGS = os.getenv(\"DATA_EMBEDDINGS\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare corpora"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load and concatenate data sets.\n",
    "\n",
    "Data sets that we ingest here need to have the following properties:\n",
    "- `identifier`: Unique identifier of the document. Consists of dataframe index prefixed with a signifier for the document series, e.g. `krp-`.\n",
    "- `date`: Date of creation of the document.\n",
    "- `year`: Creation year, derived from `date`.\n",
    "- `title`: Cleaned title of document.\n",
    "- `text`: Cleaned text of document.\n",
    "- `link`: Link to the document on the search portal of Staatsarchiv or zh.ch.\n",
    "- `stazh_ident`: Original identifier, derived from XML tag `ident`.\n",
    "\n",
    "Additional properties:\n",
    "- `series`: From document series KRP, RRB, GSZH or Amtsblatt. \n",
    "- `word_count`: Number of words in text, used to filter out short or empty documents.\n",
    "- `business_no`: Business number of RRBs. We don't have these for KRPs and GSZH yet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfa = pd.read_parquet(PREP_OUTPUT_KRP)\n",
    "dfb = pd.read_parquet(PREP_OUTPUT_RRB)\n",
    "dfc = pd.read_parquet(PREP_OUTPUT_OS)\n",
    "dfd = pd.read_parquet(PREP_OUTPUT_ABl)\n",
    "\n",
    "# To get more semantic signal we add titles to texts.\n",
    "dfa[\"text\"] = dfa.title + \" \" + dfa.text\n",
    "dfb[\"text\"] = dfb.title + \" \" + dfb.text\n",
    "dfc[\"text\"] = dfc.title + \" \" + dfc.text\n",
    "dfd[\"text\"] = dfd.title + \" \" + dfd.text\n",
    "dfa[\"series\"] = \"krp\"\n",
    "dfb[\"series\"] = \"rrb\"\n",
    "dfc[\"series\"] = \"os\"\n",
    "dfd[\"series\"] = \"abl\"\n",
    "\n",
    "df = pd.concat([dfa, dfb, dfc, dfd])\n",
    "df.reset_index(drop=True, inplace=True)\n",
    "del dfa, dfb, dfc, dfd\n",
    "\n",
    "# Add word count to filter out short documents in search app.\n",
    "df[\"word_count\"] = df.text.apply(lambda x: len(x.split(\" \")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Additional text cleaning for semantic search.\n",
    "# Mostly remove digits and single characters that do not add meaning.\n",
    "DIGITS = re.compile(r\"\\d+\")\n",
    "MULTIPLE_SPACES = re.compile(r\"[\\s]+\")\n",
    "SINGLE_CHARACTERS = re.compile(r\"\\s.\\s\")\n",
    "\n",
    "\n",
    "def semantic_text_cleaning(d):\n",
    "    d = re.sub(DIGITS, \" \", d)\n",
    "    d = re.sub(MULTIPLE_SPACES, \" \", d)\n",
    "    d = re.sub(r\"([!?.,;:•-]+){2,}\", r\"\\1\", d)\n",
    "    d = re.sub(SINGLE_CHARACTERS, \" \", d)\n",
    "    d = d.strip()\n",
    "    return d\n",
    "\n",
    "\n",
    "df[\"text_semantic\"] = df.text.parallel_apply(semantic_text_cleaning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 542479 entries, 0 to 542478\n",
      "Data columns (total 11 columns):\n",
      " #   Column         Non-Null Count   Dtype         \n",
      "---  ------         --------------   -----         \n",
      " 0   identifier     542479 non-null  object        \n",
      " 1   date           542479 non-null  datetime64[ns]\n",
      " 2   year           542479 non-null  int32         \n",
      " 3   title          542479 non-null  object        \n",
      " 4   text           542479 non-null  object        \n",
      " 5   link           542479 non-null  object        \n",
      " 6   stazh_ident    542479 non-null  object        \n",
      " 7   ref            542479 non-null  object        \n",
      " 8   series         542479 non-null  object        \n",
      " 9   word_count     542479 non-null  int64         \n",
      " 10  text_semantic  542479 non-null  object        \n",
      "dtypes: datetime64[ns](1), int32(1), int64(1), object(8)\n",
      "memory usage: 7.1 GB\n"
     ]
    }
   ],
   "source": [
    "df.to_parquet(DATA_OUTPUT_FULL)\n",
    "df.info(memory_usage=\"deep\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chunk documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 542479 entries, 0 to 542478\n",
      "Data columns (total 10 columns):\n",
      " #   Column       Non-Null Count   Dtype         \n",
      "---  ------       --------------   -----         \n",
      " 0   identifier   542479 non-null  object        \n",
      " 1   date         542479 non-null  datetime64[ns]\n",
      " 2   year         542479 non-null  int32         \n",
      " 3   title        542479 non-null  object        \n",
      " 4   text         542479 non-null  object        \n",
      " 5   link         542479 non-null  object        \n",
      " 6   stazh_ident  542479 non-null  object        \n",
      " 7   ref          542479 non-null  object        \n",
      " 8   series       542479 non-null  category      \n",
      " 9   word_count   542479 non-null  int64         \n",
      "dtypes: category(1), datetime64[ns](1), int32(1), int64(1), object(6)\n",
      "memory usage: 2.3 GB\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_parquet(DATA_OUTPUT_FULL)\n",
    "df.series = df.series.astype(\"category\")\n",
    "df.drop(columns=[\"text_semantic\"], inplace=True)\n",
    "df.info(memory_usage=\"deep\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = df.sample(frac=1).parallel_apply(\n",
    "    chunk_text, max_token_count=500, overlap_tokens=100, axis=1\n",
    ")\n",
    "df_chunks = pd.DataFrame(\n",
    "    [y for x in results.tolist() for y in x], columns=[\"identifier\", \"chunk_text\"]\n",
    ")\n",
    "\n",
    "df_chunks = pd.merge(\n",
    "    df.drop(columns=[\"text\"]), df_chunks, left_on=\"identifier\", right_on=\"identifier\"\n",
    ")\n",
    "\n",
    "df_chunks.info(memory_usage=\"deep\")\n",
    "df_chunks.to_parquet(DATA_OUTPUT_CHUNKS)"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}