synthesis_action_retriever_model_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from pprint import pprint\n",
    "import zipfile\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import random as rnd\n",
    "import pandas as pd\n",
    "import regex as re\n",
    "\n",
    "from progressbar import ProgressBar"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize required packages and input relevant models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "print(\"Spacy version\", spacy.__version__)\n",
    "lang_model = 'en_core_web_sm'\n",
    "nlp = spacy.load(lang_model)\n",
    "\n",
    "#uncomment line below to accelerate training with GPU\n",
    "#spacy.require_gpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gensim\n",
    "from gensim.models import Word2Vec\n",
    "#from gensim.models.fasttext import FastText\n",
    "print(\"Gensim verion:\", gensim.__version__)\n",
    "\n",
    "#Please provide your own word embedding\n",
    "embeddings = Word2Vec.load(\"path_to_embedding_model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.layers import Input, Dense, Activation, SimpleRNN, Bidirectional, Dropout\n",
    "\n",
    "print(\"TF version:\", tf.__version__)\n",
    "print(\"Keras version\", keras.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from synthesis_action_retriever.utils import make_spacy_tokens\n",
    "from text_cleanup import TextCleanUp\n",
    "tc = TextCleanUp()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload annotated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_dataset = './data/synthesis_action_annotated_dataset_2021-10-17.json'\n",
    "\n",
    "with open(path_to_dataset, 'r') as fp:\n",
    "    annotated_data = json.load(fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Number of annotated sentences: \", len(annotated_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_tok_thresh = 5\n",
    "max_tok_thresh = 50\n",
    "all_sentences = [\n",
    "    s for s in annotated_data \n",
    "    if len(s[\"annotations\"]) > min_tok_thresh and len(s[\"annotations\"]) < max_tok_thresh\n",
    "]\n",
    "print(\"Number of sentences for training after thresholding tokens: \", len(all_sentences))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "elements_1 = ['H', 'B', 'C', 'N', 'O', 'F', 'P', 'S', 'K', 'V', 'Y', 'I', 'W', 'U']\n",
    "elements_2 = ['He', 'Li', 'Be', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'Cl', 'Ar', 'Ca', 'Sc', 'Ti', 'Cr',\n",
    "              'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr', 'Rb', 'Sr',\n",
    "              'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In', 'Sn', 'Sb', 'Te', 'Xe',\n",
    "              'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd', 'Tb', 'Dy', 'Ho', 'Er',\n",
    "              'Tm', 'Yb', 'Lu', 'Hf', 'Ta', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi',\n",
    "              'Po', 'At', 'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'Np', 'Pu', 'Am', 'Cm', 'Bk', 'Cf',\n",
    "              'Es', 'Fm', 'Md', 'No', 'Lr', 'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt', 'Ds', 'Rg', 'Cn', 'Fl', 'Lv']\n",
    "num_set = set(\"0987654321+-()[]\")\n",
    "\n",
    "def is_formula_like(tok):\n",
    "    if all(c.islower() and not c.isdigit() for c in tok[1:]):\n",
    "        return False\n",
    "    \n",
    "    token_subs = tok\n",
    "    for el in elements_2:\n",
    "        token_subs = token_subs.replace(el, \"\")\n",
    "    for el in elements_1:\n",
    "        token_subs = token_subs.replace(el, \"\")\n",
    "    if len(token_subs) < len(tok):\n",
    "        return True\n",
    "    \n",
    "    return False\n",
    "\n",
    "def is_num_like(tok):\n",
    "    if len([c for c in tok if c in num_set])/len(tok) > 0.5:\n",
    "        return True\n",
    "    \n",
    "    if all(not c.isalpha() for c in tok): \n",
    "        return True\n",
    "    if tok[0].isdigit() and tok.islower():\n",
    "        return True\n",
    "\n",
    "    return False\n",
    "\n",
    "def replace_token_upd(tok, mode):\n",
    "    if is_num_like(tok.text) and len(tok.text) > 1:\n",
    "        return '<num>'\n",
    "    \n",
    "    if is_formula_like(tok.text) and len(tok.text) > 1:\n",
    "        return '<chem>'\n",
    "    \n",
    "    if mode == 'lemma':\n",
    "        return tok.lemma_\n",
    "#         if lemmas_freq[tok.lemma_] < 2:\n",
    "#             return '<unk>'\n",
    "#         else:\n",
    "#             return tok.lemma_\n",
    "    else:\n",
    "        return tok.text.lower()\n",
    "#         if words_freq[tok.text.lower()] < 2:\n",
    "#             return '<unk>'\n",
    "#         else:\n",
    "#             return tok.text.lower()\n",
    "\n",
    "    return tok.text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "action2num = {\n",
    "    \"\": 0,\n",
    "    'Non-altering': 0,\n",
    "    'Starting': 1,\n",
    "    'Mixing': 2,\n",
    "    'Purification': 3,\n",
    "    'Heating': 4,\n",
    "    'Shaping': 5,\n",
    "    'Cooling': 6,\n",
    "    'Reaction': 7\n",
    "}\n",
    "\n",
    "num2action = {\n",
    "    0: \"\",\n",
    "    1: 'Starting',\n",
    "    2: 'Mixing',\n",
    "    3: 'Purification',\n",
    "    4: 'Heating',\n",
    "    5: 'Shaping',\n",
    "    6: 'Cooling',\n",
    "    7: 'Reaction'\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Featurize tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stop_list = ['the', 'a', 'an', 'oftentimes', 'however', 'moreover', 'therefore', 'whereas', 'whereby', 'hence', \n",
    "             'thus', 'where']\n",
    "\n",
    "def get_embeddings(word_, embed_model):\n",
    "    \n",
    "    if word_ in [\"<start>\", \"<end>\"]:\n",
    "        return np.zeros(embed_model.trainables.layer1_size, dtype=float)\n",
    "    \n",
    "    word = tc.cleanup_text(word_).lower()\n",
    "    if word in embed_model.wv.vocab:\n",
    "        return embed_model.wv.__getitem__(word)\n",
    "    else:\n",
    "        return embed_model.wv.__getitem__(\"<unk>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_frac = 0.8\n",
    "\n",
    "num_classes = len(num2action)\n",
    "featurized_sentences = []\n",
    "\n",
    "training_size = int(train_frac*len(all_sentences))\n",
    "print(\"Training size:\", training_size)\n",
    "print(\"Test size:\", len(all_sentences) - training_size)\n",
    "\n",
    "rnd.shuffle(all_sentences)\n",
    "\n",
    "test_sents = all_sentences[training_size:]\n",
    "training_sents = all_sentences[0: training_size]\n",
    "\n",
    "bar = ProgressBar(max_value=len(training_sents))\n",
    "\n",
    "for num, sentence in enumerate(training_sents):\n",
    "    sentence_features = []\n",
    "    sentence_labels = []\n",
    "\n",
    "    spacy_tokens = spacy.tokens.Doc(nlp.vocab, words = [a[\"token\"] for a in sentence[\"annotations\"]])\n",
    "    \n",
    "    sentence_features.append(get_embeddings(\"<start>\", embeddings))\n",
    "    sentence_labels.append(np.zeros(num_classes))\n",
    "    \n",
    "    for word, annot in zip(spacy_tokens, sentence[\"annotations\"]):\n",
    "        embed_vec = get_embeddings(replace_token_upd(word, mode=\"\"), embeddings)\n",
    "        action_vec = keras.utils.to_categorical(action2num[annot[\"tag\"]], num_classes)\n",
    "        \n",
    "        sentence_features.append(embed_vec)\n",
    "        sentence_labels.append(action_vec)\n",
    "\n",
    "    sentence_features.append(get_embeddings(\"<end>\", embeddings))\n",
    "    sentence_labels.append(np.zeros(num_classes))\n",
    "\n",
    "    featurized_sentences.append(dict(\n",
    "            data = sentence_features,\n",
    "            labels = sentence_labels\n",
    "        ))\n",
    "    \n",
    "    bar.update(num)\n",
    "\n",
    "print(len(featurized_sentences))\n",
    "print(len(featurized_sentences[0]['data'][0]))  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_word_dim = embeddings.trainables.layer1_size\n",
    "seq_len = max([len(d[\"data\"]) for d in featurized_sentences])\n",
    "output_dim = num_classes\n",
    "\n",
    "print(\"Input word dimention:\", input_word_dim)\n",
    "print(\"Input sequence length:\", seq_len)\n",
    "print(\"Output dimention:\", output_dim)\n",
    "print(\"Output sequence length (same as input):\", seq_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_sentences_data = np.zeros((len(featurized_sentences), seq_len, input_word_dim), dtype='float32')\n",
    "output_tags_data = np.zeros((len(featurized_sentences), seq_len, output_dim), dtype='float32')\n",
    "\n",
    "for i, data in enumerate(featurized_sentences[0:training_size]):\n",
    "    for t, (word, tag) in enumerate(zip(data[\"data\"], data[\"labels\"])):\n",
    "        input_sentences_data[i, t] = word\n",
    "        output_tags_data[i, t] = tag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(input_sentences_data[0].shape)\n",
    "print(output_tags_data[0].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_dim = 32\n",
    "\n",
    "model = None\n",
    "X = Input(shape=(None, input_word_dim))\n",
    "#lstm = SimpleRNN(latent_dim, return_sequences=True)(X)\n",
    "lstm = Bidirectional(SimpleRNN(latent_dim, return_sequences=True, dropout=0.2, recurrent_dropout=0.2))(X)\n",
    "dense = Dense(output_dim)(lstm)\n",
    "prediction = Activation(\"softmax\")(dense)\n",
    "model = Model(inputs=X, outputs=prediction)\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "batch_size = 128\n",
    "epochs = 64\n",
    "model.compile(optimizer='rmsprop', loss='categorical_crossentropy',\n",
    "              metrics=['accuracy'])\n",
    "model.fit([input_sentences_data], output_tags_data,\n",
    "          batch_size=batch_size,\n",
    "          epochs=epochs,\n",
    "          validation_split=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = rnd.randint(0, len(all_sentences)-1)\n",
    "test_sent = all_sentences[idx]\n",
    "\n",
    "test_words = [a[\"token\"] for a in test_sent[\"annotations\"]]\n",
    "test_tags = [a[\"tag\"] for a in test_sent[\"annotations\"]]\n",
    "\n",
    "print(test_words)\n",
    "print([num2action[action2num[t]] for t in test_tags])\n",
    "\n",
    "spacy_tokens = spacy.tokens.Doc(nlp.vocab, words = test_words)\n",
    "\n",
    "input_sentences_data = np.zeros((1, seq_len, input_word_dim), dtype='float32')\n",
    "for t, (word, tag) in enumerate(zip(spacy_tokens, test_tags)):\n",
    "    embed_vec = get_embeddings(replace_token_upd(word, mode=\"\"), embeddings)\n",
    "    input_sentences_data[0, t] = embed_vec\n",
    "    \n",
    "result = model.predict(input_sentences_data)[0]\n",
    "tags_predicted = []\n",
    "for word, pred_vec in zip(test_words, result):\n",
    "    tags_predicted.append(num2action[np.argmax(pred_vec)])\n",
    "\n",
    "print(tags_predicted)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "correct = []\n",
    "missing = []\n",
    "extra = []\n",
    "wrong_tag = []\n",
    "\n",
    "tn = []\n",
    "tp = []\n",
    "fp = []\n",
    "fn = []\n",
    "\n",
    "\n",
    "bar = ProgressBar(max_value = len(test_sents))\n",
    "for sentence in test_sents:\n",
    "    \n",
    "    words = [a[\"token\"] for a in sentence[\"annotations\"]]\n",
    "    tags = [a[\"tag\"] for a in sentence[\"annotations\"]]\n",
    "    \n",
    "    correct_tags = []\n",
    "    for t in tags:\n",
    "        t = \"Mixing\" if \"Mixing\" in t else t\n",
    "        t = \"\" if t == \"Miscellaneous\" else t\n",
    "        correct_tags.append(t)\n",
    "    \n",
    "    spacy_tokens = spacy.tokens.Doc(nlp.vocab, words = words)\n",
    "\n",
    "    seq_len = len(words)\n",
    "    input_sentences_data = np.zeros((1, seq_len, input_word_dim), dtype='float32')\n",
    "    input_sentences_data[0, 0] = get_embeddings(\"<start>\", embeddings)\n",
    "    for t, word in enumerate(spacy_tokens):\n",
    "        embed_vec = get_embeddings(replace_token_upd(word, mode=\"\"), embeddings)\n",
    "        input_sentences_data[0, t] = embed_vec\n",
    "    input_sentences_data[0, -1] = get_embeddings(\"<end>\", embeddings)\n",
    "\n",
    "    result = model.predict(input_sentences_data)[0]\n",
    "    tags_predicted = [num2action[np.argmax(v)] for v in result]#[0:len(spacy_tokens)]\n",
    "    \n",
    "    sentence[\"prediction\"] = tags_predicted#[1:-1]\n",
    "    sentence[\"correct\"] = correct_tags\n",
    "\n",
    "    if tags_predicted == correct_tags:\n",
    "        correct.append(sentence)\n",
    "        if \"\".join([t for t in correct_tags+tags_predicted]) == \"\":\n",
    "            tn.append(sentence)\n",
    "        else:\n",
    "            tp.append(sentence)\n",
    "    elif len([t for t in tags_predicted if t != \"\"]) > len([t for t in correct_tags if t != \"\"]):\n",
    "        extra.append(sentence)\n",
    "        fp.append(sentence)\n",
    "    elif len([t for t in tags_predicted if t != \"\"]) < len([t for t in correct_tags if t != \"\"]):\n",
    "        missing.append(sentence)\n",
    "        fn.append(sentence)\n",
    "    else:\n",
    "        wrong_tag.append(sentence)\n",
    "    \n",
    "print(\"Correct:\", len(correct))\n",
    "print(\"Extra:\", len(extra))\n",
    "print(\"Missing:\", len(missing))\n",
    "print(\"Wrong:\", len(wrong_tag))\n",
    "print(\"Test set:\", len(test_sents))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prec = len(tp)/(len(tp)+len(fp))\n",
    "recall = len(tp)/(len(tp)+len(fn))\n",
    "accuracy = (len(tp) + len(tn))/(len(tp)+len(tn)+len(fp)+len(fn))\n",
    "f1 = 2.0*prec*recall/(prec + recall)\n",
    "\n",
    "print(\"Precision:\", round(prec, 2))\n",
    "print(\"Recall:\", round(recall, 2))\n",
    "print(\"Accuracy:\", round(accuracy, 2))\n",
    "print(\"F1:\", round(f1, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "timestr = time.strftime(\"%Y%m%d-%H%M%S\")\n",
    "if lang_model=='en_core_web_trf':\n",
    "    tf.saved_model.save(model, './output/Bi-RNN_cl7_ed100_TF_{}'.format(timestr))\n",
    "else:\n",
    "    model.save(\"./output/Bi-RNN_cl7_ed100_{}\".format(timestr))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}