examples/AnnoMI/anno_mi-client-swnu-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a SWNU-Network model\n",
    "\n",
    "In this notebook, we will go through some helper functions in the library to train and evalaute a SWNU-Network model. We will also look at the SWNU itself to see how we can use that and plug that into other PyTorch networks.\n",
    "\n",
    "Note that using the SW-Attention (SWMHAU) unit is very similar and have the same inputs as the SWNU (besides the initialisation of the unit of course). In particular, you would use the same `obtain_SWNUNetwork_input` to construct the path to input into the unit and SWMHAU-Network model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a983a5e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "from load_anno_mi import (\n",
    "    anno_mi,\n",
    "    client_index,\n",
    "    client_transcript_id,\n",
    "    output_dim_client,\n",
    "    y_data_client,\n",
    "    label_to_id_client,\n",
    "    id_to_label_client,\n",
    ")\n",
    "\n",
    "from sig_networks.scripts.swnu_network_functions import (\n",
    "    obtain_SWNUNetwork_input,\n",
    "    implement_swnu_network,\n",
    ")\n",
    "\n",
    "from sig_networks.swnu import SWNU\n",
    "\n",
    "import signatory"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f7409a03",
   "metadata": {},
   "source": [
    "## AnnoMI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f00bb922",
   "metadata": {},
   "outputs": [],
   "source": [
    "anno_mi.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below are some statistics regarding how many therapist and client utterances there are (after removing duplicates - see the [load_anno_mi.py](load_anno_mi.py) script) and the proportion of each class in the \"main_therapist_behaviour\" and \"client_talk_type\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "720d820f",
   "metadata": {},
   "outputs": [],
   "source": [
    "anno_mi[\"interlocutor\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84d5594d",
   "metadata": {},
   "outputs": [],
   "source": [
    "anno_mi[\"main_therapist_behaviour\"].value_counts() / anno_mi[\n",
    "    \"interlocutor\"\n",
    "].value_counts()[\"therapist\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44217a68",
   "metadata": {},
   "outputs": [],
   "source": [
    "anno_mi[\"client_talk_type\"].value_counts() / anno_mi[\"interlocutor\"].value_counts()[\n",
    "    \"client\"\n",
    "]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d31003be",
   "metadata": {},
   "source": [
    "## Client talk type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5a752a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(client_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e478c216",
   "metadata": {},
   "outputs": [],
   "source": [
    "label_to_id_client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f716f8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "id_to_label_client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30c4b24d",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dim_client"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "07a35e43",
   "metadata": {},
   "source": [
    "## Obtaining SBERT Embeddings\n",
    "\n",
    "We can use the `SentenceEncoder` class within `nlpsig` to obtain sentence embeddings from a model. This class uses the [`sentence-transformer`](https://www.sbert.net/docs/package_reference/SentenceTransformer.html) package and here, we have use the pre-trained `all-MiniLM-L12-v2` model - alternative models can be found [here](https://www.sbert.net/docs/pretrained_models.html).\n",
    "\n",
    "All of this is done in the `sbert_embeddings.py` script:\n",
    "- We use the `nlpsig.SentenceEncoder` class where we pass in our dataframe, `df=anno_mi`, along with the column name which contains our text, `feature_name=\"utterance_text\"`. Further, we can pass in the pre-trained sentence-transformer model we want to use `model_name=\"all-MiniLM-L12-v2\"` (see what other pre-trained models are available in the `sentence-transformer` library [here](https://www.sbert.net/docs/pretrained_models.html).\n",
    "- We can then simply obtain embeddings for each item in our dataframe by:\n",
    "    1. Loading the pre-trained model using `nlpsig.SentenceEncoder.load_pretrained_model()`\n",
    "    2. Obtain SBERT embeddings using `nlpsig.SentenceEncoder.obtain_embeddings()`\n",
    "- We lastly pickle the embeddings to `\"anno_mi_sbert.pkl\"`\n",
    "\n",
    "We will leave the next cell commented out as we have already ran this, but you will need to run this first if this does not exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89fd2dd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %run ../sbert_embeddings.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can load in the embeddings, and we see that for each item in our dataframe, we now have a corresponding SBERT representation for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87e3568e",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"anno_mi_sbert.pkl\", \"rb\") as f:\n",
    "    sbert_embeddings = pickle.load(f)\n",
    "\n",
    "sbert_embeddings.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anno_mi.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constructing paths and input for SWNU\n",
    "\n",
    "For using the Signature Window Network Unit (SWNU), for each data point, we need to construct a path of it's history. In this particular example, this is the path consisting of the SBERT representations of the conversation up to the current utterance/data point. \n",
    "\n",
    "For all path construction functionality, we can utilise the `nlpsig` library and in particular the `nlpsig.PrepareData` class. To make this even easier, there are helper functions within the `sig_networks` library: specifically the `obtain_SWNUNetwork_input` function from `sig_networks.scripts.swnu_network_functions`.\n",
    "\n",
    "For this function, we need to pass in:\n",
    "- the `method` for dimension reduction we wish to use (for full options, see `nlpsig.DimReduce`)\n",
    "- the `dimension` to reduce down to\n",
    "- the dataframe, `df`, that contains our data\n",
    "- the column name which defines the grouping of our texts, `id_column`\n",
    "    - For this example, this is the transcript/conversation ID, `id_column=\"transcript_id\"`\n",
    "- the label column `label_column`\n",
    "    - For this example, this is the client talk type, `label_column=\"client_talk_type\"`\n",
    "- the embeddings for each item in the dataframe which we obtained earlier\n",
    "- the history length to consider `k`\n",
    "\n",
    "Optionally, there may be some additional features present that we wish to add in such as any time variables or any other non-textual information. We can pass these into the `features` argument and we also have functionality to either include these features in either the path, the inputs that we concatentate at the end of SWNUNetwork, or both. Here, we will just add some time features to the path.\n",
    "\n",
    "Note that after constructing a `nlpsig.PrepareData` object, it will create some time variables such as:\n",
    "- the `timeline_index` which is simply an ordered index for each grouping given by the `id_column`\n",
    "- (if `datetime` is a column in the dataframe) the `time_encoding` which is the date of the data point as a fraction of the year, e.g. 31/01/2014 at 00:00:00 is 31/365 into the year so will be converted to 2014.0849\n",
    "- (if `datetime` is a column in the dataframe) the `time_encoding_minute` which is the time of the data point (ignoring the date) as a fraction of a minute, e.g. 00:01:30 would be converted to 1.50\n",
    "\n",
    "See `nlpsig.PrepareData` for all options implemented. Note that we could also add any other feature that is present in the dataframe (as long as is numerical - categorical ones could be converted to numerical). \n",
    "\n",
    "Here we will just use `timeline_index` and `time_encoding_minute` which refers to the index of the current data point within its conversation and the time of the utterance as fraction of a minute. We apply no standardisation to these but standardisation methods are available to be applied via the `standardise_method` argument.\n",
    "\n",
    "Lastly, since in this example, we only want to predict on the client talk type but include the therapist utterances in the history, we can slice the input by passing the indices of interest into the `path_indices` argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "swnu_network_input = obtain_SWNUNetwork_input(\n",
    "    method=\"umap\",\n",
    "    dimension=15,\n",
    "    df=anno_mi,\n",
    "    id_column=\"transcript_id\",\n",
    "    label_column=\"client_talk_type\",\n",
    "    embeddings=sbert_embeddings,\n",
    "    k=5,\n",
    "    features=[\"time_encoding_minute\", \"timeline_index\"],\n",
    "    include_features_in_path=True,\n",
    "    include_features_in_input=False,\n",
    "    seed=42,\n",
    "    path_indices=client_index,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To see how the `nlpsig` is used to construct the input, please see the source code for the `obtain_SWNUNetwork_input` function. Within this function, we have made some certain choices on how we should construct a path, but it is possible to construct the input differently with some customisation to this function and to adapt this approach for what your task requires."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(swnu_network_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "swnu_network_input.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[type(swnu_network_input[key]) for key in swnu_network_input.keys()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result of this is a dictionary with keys:\n",
    "- `x_data`: this in itself is another dictionary with keys `path` which is a three-dimensional tensor containing the batch of streams which get processed by the SWNU and `features` which is a two-dimensional tensor containing the batch of embeddings that we concatenate to the output of SWNU in the network\n",
    "- `input_channels`: this is a integer which is computed by the dimension of the dimension reduced embeddings + the number of features we want to add (15+2 in our case)\n",
    "- `embedding_dim`: this is the dimension of the embeddings we want to concatentate with (384 in our case)\n",
    "- `num_features`: this is the number of features that we're concatenating in the input (0 since we set `include_features_in_input=False`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "swnu_network_input[\"x_data\"].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "swnu_network_input[\"x_data\"][\"path\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "swnu_network_input[\"x_data\"][\"features\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    swnu_network_input[\"input_channels\"],\n",
    "    swnu_network_input[\"embedding_dim\"],\n",
    "    swnu_network_input[\"num_features\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the SWNU PyTorch class\n",
    "\n",
    "The SWNU class can be easily initialised plugged into any PyTorch network. We simply just need to initialise a SWNU object and then in a forward pass, the SWNU takes in a three dimensional tensor of batches of streams just as we've constructed here and return a fixed feature representation for each item in the batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialise a SWNU object\n",
    "swnu = SWNU(\n",
    "    input_channels=swnu_network_input[\"input_channels\"],\n",
    "    hidden_dim=5,\n",
    "    log_signature=True,\n",
    "    sig_depth=3,\n",
    "    pooling=\"signature\",\n",
    "    BiLSTM=True,\n",
    ")\n",
    "\n",
    "# pass the path through the SWNU\n",
    "features = swnu(swnu_network_input[\"x_data\"][\"path\"].float())\n",
    "\n",
    "# features is a two-dimensional tensor of shape [batch, signature_channels]\n",
    "features.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note here that the `signature_channels` here in this case is the number log-signature terms (since `log_signature=True`) for a path which has 10 channels (as we take a signature on the LSTM hidden dims using \"signature\" pooling, and taking signature depth 3:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "signatory.logsignature_channels(5, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have the input for the SWNU, we can use the `implement_swnu_network` helper function provided in the library in `sig_networks.scripts.swnu_network_functions`. Here we set up some of the arguments to pass into the function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kwargs = {\n",
    "    \"num_epochs\": 5,\n",
    "    \"x_data\": swnu_network_input[\"x_data\"],\n",
    "    \"y_data\": y_data_client,\n",
    "    \"input_channels\": swnu_network_input[\"input_channels\"],\n",
    "    \"num_features\": swnu_network_input[\"num_features\"],\n",
    "    \"embedding_dim\": swnu_network_input[\"embedding_dim\"],\n",
    "    \"log_signature\": True,\n",
    "    \"sig_depth\": 3,\n",
    "    \"pooling\": \"signature\",\n",
    "    \"swnu_hidden_dim\": 5,\n",
    "    \"ffn_hidden_dim\": [32, 32],\n",
    "    \"output_dim\": output_dim_client,\n",
    "    \"BiLSTM\": True,\n",
    "    \"dropout_rate\": 0.1,\n",
    "    \"learning_rate\": 3e-4,\n",
    "    \"seed\": 0,\n",
    "    \"loss\": \"focal\",\n",
    "    \"gamma\": 2,\n",
    "    \"split_ids\": client_transcript_id,\n",
    "    \"k_fold\": True,\n",
    "    \"patience\": 3,\n",
    "    \"verbose_results\": True,\n",
    "    \"verbose_model\": True,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: we only do a small number of epochs and try to reduce the number of parameters in this model than what we usually run so that this notebook runs a bit quicker. For better performance, we will need to train for longer and to increase the size of the model.\n",
    "\n",
    "Here, we set up most of the arguments to the SWNU-Network model as well as some training arguments such as the number of epochs, the learning rate, the loss function. We are going to do a $K$-fold analysis over a deafult of $5$ splits.\n",
    "\n",
    "Since we're doing $K$-fold analysis, we get a randomly initialised SWNU-Network model (if we set `k_fold=False`, then we would get a trained model trained on the the training split using a validation set), and a results dataframe.\n",
    "\n",
    "The dataset gets split and since we the data points in the dataset are grouped by the transcript ID, we must pass in the transcript IDs so that we have no data contamination between folds and data splits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "swnu_network, results_df = implement_swnu_network(**kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The library also has a function to do some hyperparameter searching - for SWNU, see the `swnu_network_hyperparameter_search` function in `sig_networks.scripts.swnu_network_functions`. For each of the other models, there will be similar functions that we've seen in this notebook. \n",
    "\n",
    "For examples of this using the hyperparameter search functions, see the scripts for running the experiments in the `examples/` folder in the repo."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nlpsig-networks",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}