03_prepare_datasets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import re"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prepare data for training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import generated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "not_augmented = pd.read_csv('data/dataset_not_augmented.csv')\n",
    "augmented_2x = pd.read_csv('data/dataset_augmented_2x.csv')\n",
    "augmented_5x = pd.read_csv('data/dataset_augmented_5x.csv')\n",
    "scrambled_5x = pd.read_csv('data/dataset_scrambled_5x.csv')\n",
    "augmented_10x = pd.read_csv('data/dataset_augmented_10x.csv')\n",
    "augmented_20x = pd.read_csv('data/dataset_augmented_20x.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tokenize the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def smi_tokenizer(smi: str) -> str:\n",
    "        \"\"\"\n",
    "        Tokenize a SMILES molecule or reaction.\n",
    "        \"\"\"\n",
    "        pattern =  \"(\\[[^\\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\\(|\\)|\\.|=|#|-|\\+|\\\\\\\\|\\/|:|~|@|\\?|>|\\*|\\!|\\$|\\%[0-9]{2}|[0-9])\"\n",
    "        regex = re.compile(pattern)\n",
    "        tokens = [token for token in regex.findall(smi)]\n",
    "        assert smi == ''.join(tokens)\n",
    "        return ' '.join(tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "not_augmented['source'] = not_augmented['source'].apply(smi_tokenizer)\n",
    "augmented_2x['source'] = augmented_2x['source'].apply(smi_tokenizer)\n",
    "augmented_5x['source'] = augmented_5x['source'].apply(smi_tokenizer)\n",
    "scrambled_5x['source'] = scrambled_5x['source'].apply(smi_tokenizer)\n",
    "augmented_10x['source'] = augmented_10x['source'].apply(smi_tokenizer)\n",
    "augmented_20x['source'] = augmented_20x['source'].apply(smi_tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "not_augmented['target'] = not_augmented['target'].apply(smi_tokenizer)\n",
    "augmented_2x['target'] = augmented_2x['target'].apply(smi_tokenizer)\n",
    "augmented_5x['target'] = augmented_5x['target'].apply(smi_tokenizer)\n",
    "scrambled_5x['target'] = scrambled_5x['target'].apply(smi_tokenizer)\n",
    "augmented_10x['target'] = augmented_10x['target'].apply(smi_tokenizer)\n",
    "augmented_20x['target'] = augmented_20x['target'].apply(smi_tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Export data for training with OpenNMT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def export_to_opennmt(df, path):\n",
    "    \n",
    "    train = df[df['split'] == 'train']\n",
    "    test = df[df['split'] == 'test']\n",
    "    val = df[df['split'] == 'validation']\n",
    "\n",
    "    train['source'].to_csv(path + '/src-train.txt', index=False, header=False)\n",
    "    train['target'].to_csv(path + '/tgt-train.txt', index=False, header=False)\n",
    "    test['source'].to_csv(path + '/src-test.txt', index=False, header=False)\n",
    "    test['target'].to_csv(path + '/tgt-test.txt', index=False, header=False)\n",
    "    val['source'].to_csv(path + '/src-val.txt', index=False, header=False)\n",
    "    val['target'].to_csv(path + '/tgt-val.txt', index=False, header=False)\n",
    "\n",
    "export_to_opennmt(not_augmented, 'data/opennmt/not_augmented')\n",
    "export_to_opennmt(augmented_2x, 'data/opennmt/augmented_2x')\n",
    "export_to_opennmt(augmented_5x, 'data/opennmt/augmented_5x')\n",
    "export_to_opennmt(scrambled_5x, 'data/opennmt/scrambled_5x')\n",
    "export_to_opennmt(augmented_10x, 'data/opennmt/augmented_10x')\n",
    "export_to_opennmt(augmented_20x, 'data/opennmt/augmented_20x')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chiralpred",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}