samples/ner_sample/notebooks/flair_ner.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Named Entity Recognition using Flair on CONLL-2003\n",
    "## Experiment description\n",
    "This notebook contains a ML fabric flow for Named Entity Recognition\n",
    "using the [flair NLP package](https://github.com/flairNLP/flair/)\n",
    "\n",
    "Note: This example implements all the required data objects. For a clean notebook which uses the objects already implemented in the `ner_sample` Python package [click here](flair_ner_clean.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Jupyter helpers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from typing import List, Tuple\n",
    "import copy\n",
    "\n",
    "import flair\n",
    "import requests\n",
    "import torch\n",
    "from flair.data import Corpus, Sentence\n",
    "from flair.datasets import CONLL_03\n",
    "from flair.embeddings import TokenEmbeddings, WordEmbeddings, StackedEmbeddings, PooledFlairEmbeddings\n",
    "from flair.models import SequenceTagger\n",
    "from flair.trainers import ModelTrainer\n",
    "from seqeval.metrics import f1_score, accuracy_score\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from ner_sample import ExperimentRunner\n",
    "from ner_sample.data import DataLoader\n",
    "from ner_sample.evaluation import Evaluator, EvaluationMetrics\n",
    "from ner_sample.experimentation import MlflowExperimentation\n",
    "from ner_sample.models import BaseModel\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Dataset loading\n",
    "\n",
    "First thing we do is to implement a `DataLoader`. A `DataLoader` defines the logic for obtaining a dataset. It could either fetch a dataset from a local folder, or from a remote location like the web, S3, Blob storage or similar. \n",
    "\n",
    "To implement a `DataLoader`, there are two main functions to be created:\n",
    "- `download_dataset`: A function for downloading the dataset into the local machine (should be implmented in a way that only downloads once and then checks if the dataset already exists locally)\n",
    "- `get_dataest`: A function for getting a dataset for modeling, into the experiment code itself.\n",
    "\n",
    "Note:\n",
    "- Each dataset should have a name and a version, which will be used to know exactly what data was used for this experiment. This would provide us with the possibility of reproducing the experiment.\n",
    "- The dataset obtained should already be ready for modeling. Any train/test split should be done prior to the dataset loading. We don't want random folds created every time here because then our experiments would not be comparable (each one would use a different subset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> In this case, the dataset, *CONLL03*, will be downloaded from Github. \n",
    "> It has three folds: *eng.train* (train), *eng.testa* (dev) and *eng.testb* (test). We will be using the `flair` package to represent the dataset as a flair `Corpus`.\n",
    "\n",
    "> Additional parameters (**downsample** in this case, which samples the dataset) can be passed to the as kwargs to the super init method, and they will be automatically logged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConllDataLoader(DataLoader):\n",
    "    def __init__(\n",
    "        self,\n",
    "        dataset_name=\"conll_03\",\n",
    "        dataset_version=\"1\",\n",
    "        local_data_path=\"../data/processed/\",\n",
    "        dataset_path=\"https://raw.githubusercontent.com/glample/tagger/master/dataset/\",\n",
    "        downsample=0.05\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Data Loader for the CONLL 03 dataset.\n",
    "        download_dataset downloads the three datasets (train, testa and testb) from Github\n",
    "        get_dataset returns a flair Corpus object holding the three datasets.\n",
    "        \"\"\"\n",
    "        self.folds = (\"eng.train\", \"eng.testa\", \"eng.testb\")\n",
    "        self.local_data_path = local_data_path\n",
    "        self.dataset_path = dataset_path\n",
    "        self.downsample = downsample\n",
    "        super().__init__(dataset_name=dataset_name, \n",
    "                         dataset_version=dataset_version,\n",
    "                         downsample=downsample)\n",
    "\n",
    "    def download_dataset(self) -> None:\n",
    "        if self.dataset_name == \"conll_03\" and self.dataset_version == \"1\":\n",
    "\n",
    "            for fold in self.folds:\n",
    "                local_path = Path(self.local_data_path, self.dataset_name).resolve()\n",
    "                fold_path = self.dataset_path + fold\n",
    "                if not local_path.exists():\n",
    "                    local_path.mkdir(parents=True)\n",
    "\n",
    "                dataset_file = Path(local_path, fold)\n",
    "                if dataset_file.exists():\n",
    "                    print(\"Dataset already exists, skipping download\")\n",
    "                    return\n",
    "\n",
    "                response = requests.get(fold_path)\n",
    "                dataset_raw = response.text\n",
    "                with open(dataset_file, \"w\") as f:\n",
    "                    f.write(dataset_raw)\n",
    "                print(f\"Finished writing fold {fold} to {self.local_data_path}\")\n",
    "\n",
    "            print(\n",
    "                f\"Finished downloading dataset {self.dataset_name} version {self.dataset_version}\"\n",
    "            )\n",
    "\n",
    "        else:\n",
    "            raise ValueError(\"Selected dataset was not found\")\n",
    "\n",
    "    def get_dataset(self) -> Tuple:\n",
    "        try:\n",
    "            corpus = CONLL_03(base_path=self.local_data_path, in_memory=True)\n",
    "            corpus = corpus.downsample(self.downsample)  # Just for example purposes\n",
    "\n",
    "            train = corpus  # includes train and dev\n",
    "\n",
    "            test = corpus.test\n",
    "\n",
    "            # Copy labels to a new tag (Flair overrides the ner tag during prediction)\n",
    "            for sentence in test:\n",
    "                for token in sentence.tokens:\n",
    "                    token.annotation_layers[\"gold_ner\"] = copy.deepcopy(\n",
    "                        token.annotation_layers[\"ner\"]\n",
    "                    )\n",
    "                    token.annotation_layers[\"ner\"][0].value = \"O\"\n",
    "\n",
    "            return train, test\n",
    "\n",
    "        except FileNotFoundError:\n",
    "            print(\n",
    "                f\"Dataset {self.dataset_name} with version {self.dataset_version} not found in data/raw\"\n",
    "            )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load data\n",
    "Once we have implemented our `DataLoader`, we can just instantiate it and call `download_dataset()` and then 'get_dataset()'. This way we ensure that our notebook can be run anywhere."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "data_loader = ConllDataLoader(dataset_name = \"conll_03\", dataset_version=\"1\")\n",
    "data_loader.download_dataset()\n",
    "train_corpus, test = data_loader.get_dataset() #train_corpus is a flair Corpus containing train and dev\n",
    "train = train_corpus.train\n",
    "dev = train_corpus.dev\n",
    "\n",
    "print(f\"Train set type: {type(train)}\")\n",
    "print(f\"Test set type: {type(test)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"First sample in train sample:\\n {train.dataset[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "### Experimentation\n",
    "\n",
    "The next phase is the experiment logger definition. The default one uses MLflow, but the API is generic and can be extended to any experiment logging mechanism. \n",
    "The experimentation class is in charge of collecting all the parameters and metrics the experiments emit along the way (from the dataset name and version, through model hyperparams and up to the final metric values).\n",
    "\n",
    "To use the default one, just call `MlflowExperimentation()`\n",
    "\n",
    "> Note: If you plan to use Mlflow hosted in Databricks, follow these steps:\n",
    "1. Pass `tracking_uri='databricks'` to the `MlflowExperimentation` object\n",
    "2. See [this doc on how to create a personal access token](https://docs.databricks.com/dev-tools/api/latest/authentication.html#token-management) \n",
    "3. See [this doc on setting up databricks-cli](https://docs.microsoft.com/en-us/azure/databricks/dev-tools/cli/)\n",
    "4. [Create new experiment on Mlflow](https://docs.microsoft.com/en-us/azure/databricks/applications/mlflow/) (if needed)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experimentation = MlflowExperimentation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Modeling\n",
    "\n",
    "The next step is writing our actual model, with optional preprocessing and postprocessing.\n",
    "\n",
    "The class to implement is `BaseModel` which exposes the sklearn-style `fit` and `predict` functions that needs to be implemented.\n",
    "\n",
    "Note:\n",
    "- The model class needs to define which parameters should be logged, by adding keys and values to the self.hyper_params dict, or by passing the variables to the super constructor.\n",
    "- The base class contains fields for DataProcessors: preprocessor and postprocessor. Use these if you want the preprocessing or postprocessing to occur during the model call (which makes it easier to operationalize the model on a new environment, without having to provide all the preprocessinr and postprocessing scripts.\n",
    "- It is also possible to pass the `Experimentation` object, if it is required during training (for example, while storing values for each epoch during model training)\n",
    "\n",
    "> This example uses the `flair` framework to create a NER model. All the model hyper parameters are added as class variables and sent to the parent class constructor for logging. Some hyperparameters are collected from the actual pytorch model (in `get_hyper_params`), to have better coverage of hyperparameters for each run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class FlairNERModel(BaseModel):\n",
    "    def __init__(\n",
    "        self,\n",
    "        corpus: Corpus,\n",
    "        hidden_size: int = 256,\n",
    "        pooling: str = \"min\",\n",
    "        word_embeddings: str = \"glove\",\n",
    "        train_with_dev: bool = True,\n",
    "        max_epochs: int = 10,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        NER detector using the Flair NLP package.\n",
    "        Source: https://github.com/flairNLP/flair/blob/master/resources/docs/EXPERIMENTS.md\n",
    "        All class inputs (except for the corpus) are model hyper parameters.\n",
    "        They are then directed to the base class and get logged into the experiment logger\n",
    "        \"\"\"\n",
    "        self.tag_type = \"ner\"\n",
    "        self.tag_dictionary = None\n",
    "        self.tagger = None\n",
    "        self.embeddings = None\n",
    "\n",
    "        self.hidden_size = hidden_size\n",
    "        self.pooling = pooling\n",
    "        self.word_embeddings = word_embeddings\n",
    "        self.train_with_dev = train_with_dev\n",
    "        self.max_epochs = max_epochs\n",
    "\n",
    "        self.set_tagger_definition(corpus)\n",
    "\n",
    "        hyper_params = self.get_hyper_params(\n",
    "            hidden_size=hidden_size,\n",
    "            pooling=pooling,\n",
    "            word_embeddings=word_embeddings,\n",
    "            train_with_dev=train_with_dev,\n",
    "            max_epochs=max_epochs,\n",
    "        )\n",
    "\n",
    "        super().__init__(**hyper_params)\n",
    "\n",
    "    def fit(self, X, y=None) -> None:\n",
    "        # initialize trainer\n",
    "        trainer: ModelTrainer = ModelTrainer(self.tagger, X)\n",
    "\n",
    "        trainer.train(\n",
    "            \"models/taggers/flair-ner\",\n",
    "            train_with_dev=self.train_with_dev,\n",
    "            max_epochs=self.max_epochs,\n",
    "        )\n",
    "\n",
    "    def predict(self, X):\n",
    "        tagged_sentences = []\n",
    "        for sentence in tqdm(X):\n",
    "            self.tagger.predict(sentence)\n",
    "            tagged_sentences.append(sentence)\n",
    "        print(f\"Tagged {len(tagged_sentences)} sentences\")\n",
    "        return tagged_sentences\n",
    "\n",
    "    def get_hyper_params(self, **hyper_params):\n",
    "        basic_params = {\n",
    "            param_name: param_value\n",
    "            for (param_name, param_value) in self.tagger.__dict__.items()\n",
    "            if type(param_value) in (bool, float, int, str)\n",
    "        }\n",
    "        hyper_params.update(basic_params)\n",
    "        return hyper_params\n",
    "\n",
    "    def set_embeddings_definition(self):\n",
    "        \"\"\"\n",
    "        Sets the embedding layers used by this tagger\n",
    "        \"\"\"\n",
    "        # initialize embeddings\n",
    "        embedding_types: List[TokenEmbeddings] = [\n",
    "            # Word embeddings (default = GloVe)\n",
    "            WordEmbeddings(self.word_embeddings),\n",
    "            # contextual string embeddings, forward\n",
    "            PooledFlairEmbeddings(\"news-forward\", pooling=self.pooling),\n",
    "            # contextual string embeddings, backward\n",
    "            PooledFlairEmbeddings(\"news-backward\", pooling=self.pooling),\n",
    "        ]\n",
    "        self.embeddings: StackedEmbeddings = StackedEmbeddings(\n",
    "            embeddings=embedding_types\n",
    "        )\n",
    "\n",
    "    def set_tagger_definition(self, corpus: Corpus):\n",
    "        \"\"\"\n",
    "        Returns the definition of the Flair SequenceTagger (the full model)\n",
    "        :param corpus: Used only for setting the tag_dictionary\n",
    "        \"\"\"\n",
    "\n",
    "        if not self.embeddings:\n",
    "            self.set_embeddings_definition()\n",
    "        self.tag_dictionary = corpus.make_tag_dictionary(tag_type=self.tag_type)\n",
    "\n",
    "        tagger: SequenceTagger = SequenceTagger(\n",
    "            hidden_size=self.hidden_size,\n",
    "            embeddings=self.embeddings,\n",
    "            tag_dictionary=self.tag_dictionary,\n",
    "            tag_type=self.tag_type,\n",
    "            use_crf=False,\n",
    "        )\n",
    "        self.tagger = tagger\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Model training\n",
    "\n",
    "The model we just created can be called and fitted. Alternatively, we can postpone the fit to the last part, which performs a full experiment cycle.\n",
    "\n",
    "> In this example, we skip training as it takes a lot of time. Instead, we load a pre-trained model directly from `flair`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = FlairNERModel(corpus=train_corpus)\n",
    "\n",
    "TRAIN=False\n",
    "\n",
    "if TRAIN:\n",
    "    model.fit(corpus=train_corpus)\n",
    "else:\n",
    "    # Simulate training has finished by downloading a pretrained model\n",
    "    model.tagger = SequenceTagger.load('ner')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Prediction\n",
    "\n",
    "Once we have a fitted model, we can run the experiment to validate its performance and log results. Before that, let's verify that we get something meaningful when calling the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_sentence = Sentence(\"In Penny Lane, there is a barber showing photographs\")\n",
    "\n",
    "model.predict([example_sentence])\n",
    "for token in example_sentence.tokens:\n",
    "    print(f\" {token.text} | {token.get_tag('ner')}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks good!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Model evaluation\n",
    "In this phase, we define how the model should be evaluated. There are two main building blocks:\n",
    "- `Evaluator`: Which holds the logic for how evaluation takes place. The function to implement is `evaluate`.\n",
    "- `EvaluationMetrics`: Which holds the actual values of metrics. The function to implement is `get_metrics`.\n",
    "\n",
    "> In this example we implmement `NEREvaluator` and `NEREvaluationMetrics` with our specific logic and metrics. We use the `semeval` package to calculate **f1** and **accuracy** metrics for the NER task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class NEREvaluationMetrics(EvaluationMetrics):\n",
    "    \"\"\"\n",
    "    This class holds the metrics calculated during the experiment run\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, f1, accuracy):\n",
    "        self.f1 = f1\n",
    "        self.accuracy = accuracy\n",
    "        super().__init__()\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"F1 score: {self.f1}, Accuracy score: {self.accuracy}\"\n",
    "\n",
    "    def get_metrics(self):\n",
    "        \"\"\"\n",
    "        Return a dict with f1 and accuracy values\n",
    "        \"\"\"\n",
    "        return { \"f1\": self.f1, \"accuracy\":self.accuracy }\n",
    "\n",
    "\n",
    "class NEREvaluator(Evaluator):\n",
    "    \"\"\"\n",
    "    This class holds the logic for evaluating a prediction outcome\n",
    "    y_test in our case is None\n",
    "    \"\"\"\n",
    "\n",
    "    def evaluate(self, y_test, predictions) -> NEREvaluationMetrics:\n",
    "\n",
    "        golds = []\n",
    "        predicted = []\n",
    "        for sentence in predictions:\n",
    "            gold_tags = [token.get_tag(\"gold_ner\").value for token in sentence.tokens]\n",
    "            golds.append(gold_tags)\n",
    "            predicted_tags = [token.get_tag(\"ner\").value for token in sentence.tokens]\n",
    "            predicted.append(predicted_tags)\n",
    "\n",
    "        f1 = f1_score(golds, predicted)\n",
    "        accuracy = accuracy_score(golds, predicted)\n",
    "        return NEREvaluationMetrics(f1=f1, accuracy=accuracy)\n",
    "\n",
    "\n",
    "evaluator = NEREvaluator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Running an experiment\n",
    "\n",
    "To run the full experiment, we leverage the `ExpreimentRunner` class. This class is in charge of evaluating the model on a test dataset, calculating metrics, collecting all params and metrics and logging them to the experiment logger. It's like an experiment orchastrator. \n",
    "In additional to all the collected params and metrics, one could add additional params to the call to ExperimentRunner and these will too be collected. \n",
    "\n",
    "> In many cases the `ExperimentRunner` class could be used it without any modification, but if modifications are needed, just make sure that you implement the various functions (, and also verify that the different params and metrics are logged correctly (in the `__init__`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instantiate the `ExperimentRunner` object, while passing all the previous building blocks.\n",
    "\n",
    "Finally, we call `experiment_runner.evaluate()` to perform prediction on the supplied test set, calculate metrics and store everything in the experiment logger."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "experiment_runner = ExperimentRunner(\n",
    "    model=model,\n",
    "    X_train=train,\n",
    "    X_test=test,\n",
    "    data_loader=data_loader,\n",
    "    log_experiment=True,\n",
    "    experiment_logger=experimentation,\n",
    "    evaluator=evaluator,\n",
    "    experiment_name=\"Experiment\",\n",
    "    data_scientist=\"Omri\"\n",
    ")\n",
    "\n",
    "results = experiment_runner.evaluate()\n",
    "print(results)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary\n",
    "\n",
    "This example flow demonstrates how to use the different building blocks in this framework.\n",
    "\n",
    "**Possible next steps:**\n",
    "1. Implement the different modules in the Python package, and use them in other notebooks / scripts / modules\n",
    "2. Run `mlflow ui` from this notebook's path and observe the different parameters and metrics stored\n",
    "3. Create a notebook template for your experiment, which can be used to generate new notebooks containing the needed objects (lodaing data, experimentation, evaluation, run experiment)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start the Mlflow UI, run `!mlflow ui` and open http://localhost:5000/#/ in your browser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mlflow ui"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open http://localhost:5000/#/"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "kernelspec": {
   "display_name": "PyCharm (ner_sample)",
   "language": "python",
   "name": "pycharm-e86bbb65"
  },
  "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}