docs/tutorials/01-huggingface.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "respected-allowance",
   "metadata": {},
   "source": [
    "# 🤗 Using Rubrix to explore NLP data with Hugging Face datasets and transformers \n",
    "\n",
    "In this tutorial, we will walk through the process of using Rubrix to explore NLP datasets in combination with the amazing `datasets` and `transformer` libraries from Hugging Face.\n",
    "\n",
    "\n",
    "## Introduction\n",
    "\n",
    "**Our goal is to show you how to store and explore NLP datasets using Rubrix** for use cases like training data management or model evaluation and debugging."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "small-guarantee",
   "metadata": {},
   "source": [
    "The tutorial is organized into three parts:\n",
    "\n",
    "1. **Storing and exploring text classification data**: We will use the 🤗 `datasets` library and Rubrix to store text classification datasets.\n",
    "2. **Storing and exploring token classification data**: We will use the 🤗 `datasets` library and Rubrix to store token classification data.\n",
    "2. **Exploring predictions**: We will use a pretrained 🤗 `transformers` model and store its predictions into Rubrix to explore and evaluate our pretrained model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "strong-portfolio",
   "metadata": {},
   "source": [
    "## Install tutorial dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e6425f6",
   "metadata": {},
   "source": [
    "In this tutorial we will be using `transformers` and `datasets` libraries. If you do not have them installed, run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "purple-startup",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install torch -qqq\n",
    "%pip install transformers -qqq\n",
    "%pip install datasets -qqq\n",
    "%pip install tdqm -qqq # for progress bars"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "executed-gravity",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Setup Rubrix\n",
    "\n",
    "If you have not installed and launched Rubrix, check the [Setup and Installation guide](https://docs.rubrix.ml/en/latest/getting_started/setup%26installation.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d0e66e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "revolutionary-physics",
   "metadata": {},
   "source": [
    "## 1. Storing and exploring text classification training data\n",
    "\n",
    "Rubrix allows you to track data for different NLP tasks (such as `Token Classification` or `Text Classification`). \n",
    "\n",
    "With Rubrix you can track both training data and predictions from models. In this part, we will focus only on training data. Typically, training data is data which has been curated or annotated by a human. Other terms for this same concept are: ground-truth data, \"gold-standard\" data, or even \"annotated\" data.\n",
    "\n",
    "In this part of the tutorial, you will learn how to use 🤗 datasets library for quick exploration of `Text Classification` and `Token Classification` training data. This is useful during model development, for getting a sense of the data, identifying potential issues, debugging, etc. Here we will use rather static \"research \"datasets but Rubrix really shines when you are collecting and using training data in the wild, or in other words in real data science projects.\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "younger-intranet",
   "metadata": {},
   "source": [
    "### Text classification with the `tweet_eval` dataset (Emoji classification)\n",
    "\n",
    "Text classification deals with predicting in which categories a text fits. As if you're shown an image you could quickly tell if there's a dog or a cat in it, we build NLP models to distinguish between a Jane Austen's novel or a Charlotte Bronte's poem. It's all about feeding models with labeled examples and see how it start predicting over the very same labels.\n",
    "\n",
    "In this first case, we are going to play with `tweet_eval`, a dataset with a bunch of tweets from different authors and topics and the sentiment it transmits. This is, in fact, a very common NLP task called Sentiment Analysis, but with a cool tweak: we are representing these sentiments with emojis. Each tweet comes with a number between 0 and 19, which represents different emojis. You can see each one in a cell below or in the [tweet_eval site](https://huggingface.co/datasets/tweet_eval) at 🤗 Hub.\n",
    "\n",
    "First of all, we are going to load the dataset from 🤗 Hub and visualize its content."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "better-abortion",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"tweet_eval\", 'emoji', script_version=\"master\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "naval-category",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = dataset['train'].features['label'].names; labels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "authentic-battlefield",
   "metadata": {},
   "source": [
    "Usually, datasets are divided into train, validation and test splits, and each one of them is used in a certain part of the training. For now, we can stick to the training split, which usually contains the majority of the instances of a dataset. Let's see what's inside!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "educational-mother",
   "metadata": {},
   "outputs": [],
   "source": [
    "with dataset['train'].formatted_as(\"pandas\"):\n",
    "    print(dataset['train'][:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "suspected-healthcare",
   "metadata": {},
   "source": [
    "Now, we are going to create our records from this dataset and log them into rubrix. Rubrix comes with `TextClassificationRecord` and `TokenClassificationRecord` classes, which can be created from a dictionary. These objects passes information to rubrix  about the input of the model, the predictions obtained and the annotations made, as well as a metadata field for other important details. \n",
    "\n",
    "In our case, we haven't predicted anything, so we are only going to include the labels of each instance as annotations, as we know they are the ground truth. We will also include each tweet into inputs, and specify in the metadata section that we are into the training split. Once `records` is populated, we can log it with `rubric.log()`, specifying the name of our dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "electrical-fault",
   "metadata": {},
   "outputs": [],
   "source": [
    "records = []\n",
    "\n",
    "for record in dataset['train']:\n",
    "    records.append(rb.TextClassificationRecord(\n",
    "        inputs=record[\"text\"],\n",
    "        annotation=labels[record[\"label\"]],\n",
    "        annotation_agent=\"https://huggingface.co/datasets/tweet_eval\",\n",
    "        metadata={\"split\": \"train\"},\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ruled-paragraph",
   "metadata": {},
   "outputs": [],
   "source": [
    "rb.log(records=records, name=\"tweet_eval_emojis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fad0f407-a388-4bed-98f6-7e23368bf0b8",
   "metadata": {},
   "source": [
    "![Tweet eval dataset](img/tweet_eval_emojis.png \"Tweet eval dataset\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "metallic-montgomery",
   "metadata": {},
   "source": [
    "Thanks to our metadata section in the Text Classification Record, we can log tweets from the validation and test splits in the same dataset to explore them using the Metadata filters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "medium-crawford",
   "metadata": {},
   "outputs": [],
   "source": [
    "records_validation = []\n",
    "\n",
    "for record in dataset['validation']:\n",
    "    records_validation.append(rb.TextClassificationRecord(\n",
    "        inputs=record[\"text\"],\n",
    "        annotation=labels[record[\"label\"]],\n",
    "        annotation_agent=\"https://huggingface.co/datasets/tweet_eval\",\n",
    "        metadata={\"split\": \"validation\"},\n",
    "        )\n",
    "    )\n",
    "\n",
    "rb.log(records=records_validation, name=\"tweet_eval_emojis\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "raised-sally",
   "metadata": {},
   "outputs": [],
   "source": [
    "records_test = []\n",
    "\n",
    "for record in dataset['test']:\n",
    "    records_test.append(rb.TextClassificationRecord(\n",
    "        inputs=record[\"text\"],\n",
    "        annotation=labels[record[\"label\"]],\n",
    "        annotation_agent=\"https://huggingface.co/datasets/tweet_eval\",\n",
    "        metadata={\"split\": \"test\"},\n",
    "        )\n",
    "    )\n",
    "\n",
    "rb.log(records=records_test, name=\"tweet_eval_emojis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9769f30-a274-44d1-901c-5f9121f2e77c",
   "metadata": {},
   "source": [
    "![Tweet eval dataset](img/tweet_eval_filter.png \"Tweet eval dataset\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dutch-vitamin",
   "metadata": {},
   "source": [
    "### Natural language inference with the `MRPC` dataset\n",
    "\n",
    "Natural Language Inference (NLI) is also a very common NLP task, but a little bit different to regular Text Classification. In NLI, the model receives a premise and a hypothesis, and it must figure out if the premise hypothesis is true or not given the premise. We have three categories: entailment (true), contradiction (false) or neutral (undetermined or unrelated). With the premise *\"We live in a flat planet called Earth\"*, the hypothesis *\"The Earth is flat\"* must be classified as entailment, as it is stated in the premise. NLI works with a sort of close-world assumption, in that everything not defined in the premise cannot be suppoused from the real world.\n",
    "\n",
    "Another key difference from Text Classification is that the input come in pairs of two sentences or texts, not only one. Text Classification treats its input as a cohesive and correlated unit, while NLI treats its input as a pair and tries to find correlation. \n",
    "\n",
    "To play around with NLI we are going to use 🤗 Hub [GLUE benchmark](https://huggingface.co/datasets/glue) over the MRPC task. GLUE is a well-known benchmark resource for NLP, and allow us to use its data directly over the Microsoft Research Paraphrase Corpus, a corpus of online news.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "simplified-stationery",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "dataset = load_dataset('glue', 'mrpc', split='train') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "structural-corruption",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "peripheral-chancellor",
   "metadata": {},
   "source": [
    "We can see the two input sentences instead of one. In order to simplify the workflow, let's just test if they are equivalent or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "horizontal-surgery",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = dataset.features['label'].names ; labels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "elegant-holmes",
   "metadata": {},
   "source": [
    "Populating our record list follows the same procedure as in Text Classification, adapting our input to the new scenario of pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wrapped-portsmouth",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "records=[]\n",
    "\n",
    "for record in dataset:\n",
    "    records.append(rb.TextClassificationRecord(\n",
    "       inputs={\n",
    "           \"sentence1\": record[\"sentence1\"], \n",
    "           \"sentence2\": record[\"sentence2\"]\n",
    "        },\n",
    "        annotation=labels[record[\"label\"]],\n",
    "        annotation_agent=\"https://huggingface.co/datasets/glue#mrpc\",\n",
    "        metadata={\"split\": \"train\"},\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "virgin-honor",
   "metadata": {},
   "outputs": [],
   "source": [
    "rb.log(records=records, name=\"mrpc\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeb5b784-b2b7-4a51-84da-f3e9ce7c2b3f",
   "metadata": {},
   "source": [
    "Once your dataset is logged you can explore it using filters, keyword-based search and with [Elasticsearch's query string DSL](https://docs.rubrix.ml/en/latest/reference/rubrix_webapp_reference.html#search-input). \n",
    "\n",
    "For example, the folllowing query `inputs.sentence2:(not or dont)` lets you browse all examples containing `not` or `dont` inside the sentence2 field, which you can further filter by `Annotated as` to browse examples belonging to a specific category (e.g., `not_equivalent`)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "308dcd0d-4e64-4ba5-ae89-df9775cc0dfc",
   "metadata": {},
   "source": [
    "![MRPC dataset](img/mrpc_neg.png \"MRPC dataset\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "searching-princeton",
   "metadata": {},
   "source": [
    "### Multilabel text classification with `go_emotions` dataset\n",
    "\n",
    "Another similar task to Text Classification, but yet a bit different, is Multilabel Text Classification. Just one key difference: more than one label may be predicted. While in a regular Text Classification task we may decide that the tweet *\"I can't wait to travel to Egypts and visit the pyramids\"* fits into the hastag **#Travel**, which is accurate, in Multilabel Text Classification we can classify it as more than one hastag, like **#Travel #History #Africa #Sightseeing #Desert**.\n",
    "\n",
    "In Text Classification, the category with the highest score (which our model predicted) is going to be the category predicted, but in this task we need to establish a threshold, a value between 0 and 1, from which we will classify the labels as predictions or not. If we set it to 0.5, only categories with more than a 0.5 probability value will be considered predictions. \n",
    "\n",
    "To get used to this task and see how we can log data to Rubrix, we are going to use 🤗 Hub [go_emotions dataset](https://huggingface.co/datasets/go_emotions), with comments from different reddit forums and an associated sentiment (this experiment would also be considered Sentiment Analysis)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beautiful-champagne",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset('go_emotions', split='train[0:10]')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "legendary-battle",
   "metadata": {},
   "source": [
    "Here's an example of an instance of the datasets, and the different labels, ordered. Each label will be represented in the dataset as a number, but we will translate to its name before logging to rubrix, to see things more clearly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plastic-louis",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "persistent-antique",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = dataset.features['labels'].feature.names; labels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "floppy-penny",
   "metadata": {},
   "source": [
    "Now, we need to add a confidence value to our annotation, from 0 to 1. As these are all ground truths, we consider they have the maximum probability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "verbal-colorado",
   "metadata": {},
   "outputs": [],
   "source": [
    "records= []\n",
    "\n",
    "for record in dataset:\n",
    "    records.append(rb.TextClassificationRecord(\n",
    "        inputs={\"text\": record[\"text\"]},\n",
    "        annotation=[labels[cls] for cls in record['labels']],\n",
    "        annotation_agent=\"https://huggingface.co/datasets/go_emotions\",\n",
    "        multi_label=True,\n",
    "        metadata={\n",
    "            \"split\": \"train\"\n",
    "            },\n",
    "        )\n",
    "    )\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fantastic-hardwood",
   "metadata": {},
   "source": [
    "And logging is just as easy as before!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "unexpected-rotation",
   "metadata": {},
   "outputs": [],
   "source": [
    "rb.log(records=records, name=\"go_emotions\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "everyday-delivery",
   "metadata": {},
   "source": [
    "## 2. Storing and exploring token classification training data\n",
    "\n",
    "In this second part, we will cover Token Classification while still using 🤗 datasets library. These kind of NLP tasks aim to divide the input text into words, or syllabes, and assign certain values to them. Think about giving each word in a sentence its gramatical category, or highlight which parts of a medical report belong to a certain speciality. \n",
    "\n",
    "We are going to cover a few cases using 🤗 datasets, and see how `TokenClassificationRecord` allows us to log data in rubrix in a similar fashion."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tested-pioneer",
   "metadata": {},
   "source": [
    "### Named-Entity Recognition with `wnut17` dataset\n",
    "\n",
    "Named-Entity Recognition (NER) seeks to locate and classify named entities metioned in unstructured text into pre-defined categories. And, what's powerful about NER is that this predefined categories can be whatever we want. Maybe gramatical categories, and be the best at syntax analysis in our English class, maybe person names, or organizations, or even medical codes.\n",
    "\n",
    "For this case, we are going to use 🤗 Hub [WNUT 17 dataset](https://huggingface.co/datasets/wnut_17), about rare entities on written text. Take for example the tweet “so.. kktny in 30 mins?” - even human experts find entity kktny hard to detect and resolve. This task will evaluate the ability to detect and classify novel, emerging, singleton named entities in written text.\n",
    "\n",
    "As always, let's first dive into the data and see how it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "affected-satin",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"wnut_17\", split=\"train[0:10]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "opposite-cloud",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accurate-proportion",
   "metadata": {},
   "source": [
    "We can see a list of tags and the tokens they are refering to. We have the following rare entities in this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cultural-channels",
   "metadata": {},
   "outputs": [],
   "source": [
    "for entity, token in zip(dataset[0][\"ner_tags\"], dataset[0][\"tokens\"]):\n",
    "    if entity != 0:\n",
    "        print(f\"\"\"{token}: {dataset.features[\"ner_tags\"].feature.names[entity]}\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "lyric-continent",
   "metadata": {},
   "source": [
    "So, it make a lot of sense to translate these tags into NER tags, which are much more self-explanatory than an integer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "composed-institute",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = dataset.map(lambda instance: {\"ner_tags_translated\": [dataset.features[\"ner_tags\"].feature.names[tag] for tag in instance[\"ner_tags\"]]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "advisory-vessel",
   "metadata": {},
   "source": [
    "What we did is a mapping function over 🤗  dataset, which allow us to make changes in every instance of the dataset. The very same instance that we printed before is much more readable now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "political-theory",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "instrumental-bullet",
   "metadata": {},
   "source": [
    "Info about the meaning of the tags is available [here](https://huggingface.co/datasets/viewer/?dataset=wnut_17), but to sum up, *Empire* and *ESB* has been classified as **B-LOC**, or beggining of a location name, *State* and *Building* has been classified as **I-LOC** or intermediate/final of a location name.\n",
    "\n",
    "We need to transform a bit this information, providing an entity annotation. Entity annotations are simply tuples, with the following structure \n",
    "\n",
    "```python\n",
    "(label, start_position, end_position)\n",
    "```\n",
    "\n",
    "Let's create a function that transform our dataset records into entities. It's a bit weird, but don't worry! What's doing inside is getting the entities information as shown above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "engaging-possible",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_entities(record):\n",
    "\n",
    "    entities, text, nr_tokens = [], \" \".join(record[\"tokens\"]), len(record[\"tokens\"])\n",
    "    token_start_indexes = [text.rfind(substr) for substr in [\" \".join(record[\"tokens\"][i:]) for i in range(nr_tokens)]]\n",
    "    \n",
    "    entity = None\n",
    "    for i, tag, start in zip(range(nr_tokens), record[\"ner_tags_translated\"], token_start_indexes):\n",
    "        # end of entity\n",
    "        if entity is not None and (not tag.startswith(\"I-\") or i == nr_tokens -1):\n",
    "            entity += (start-1,)\n",
    "            entities.append(entity)\n",
    "            entity = None\n",
    "        # start new entity\n",
    "        if entity is None and tag.startswith(\"B-\"):\n",
    "            entity = (tag[2:], start) \n",
    "    \n",
    "    return entities"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "protective-manufacturer",
   "metadata": {},
   "source": [
    "Let's proceed and create a record list to log it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "weighted-heath",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "records = []\n",
    "\n",
    "for record in dataset:\n",
    "    entities = parse_entities(record)\n",
    "    records.append(rb.TokenClassificationRecord(\n",
    "        text=\" \".join(record[\"tokens\"]),\n",
    "        tokens=record[\"tokens\"],\n",
    "        annotation=entities,\n",
    "        annotation_agent=\"https://huggingface.co/datasets/wnut_17\",\n",
    "        metadata={\n",
    "            \"split\": \"train\"\n",
    "            },\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acoustic-comedy",
   "metadata": {},
   "outputs": [],
   "source": [
    "records[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "helpful-advocacy",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rb.log(records=records, name=\"ner_wnut_17\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "floppy-program",
   "metadata": {},
   "source": [
    "### Part of speech tagging with `conll2003` dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "anonymous-patrol",
   "metadata": {},
   "source": [
    "Another NLP task related to token-level classification is Part-of-Speech tagging (POS tagging). In it we will identify names, verbs, adverbs, adjectives...based on the context and the meaning of the words. It is a little bit trickier than having a huge dictionary where we can look up that *drink* is a verb and *dog* is a name. Many words change its gramatical type according to the context of the sentence, and here is where AI comes to save the day.\n",
    "\n",
    "With just our dictionary and a regular script, *dog* in `The sailor dogs the hatch.` would be classified as a name, because *dog* is a name, right? A trained NLP model would step up and say *No! That's is a very common example to ilustrate the ambiguity of words. It is a verb!*. Or maybe it would just say *verb*. That's up to you.\n",
    "\n",
    "In this [dataset](https://huggingface.co/datasets/conll2003) from 🤗 hub, we will see how differente sentence has POS and NER tags, and how we can log this POS tag information into Rubrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sonic-reflection",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"conll2003\", split=\"train[0:10]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "understood-charter",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "strange-postcard",
   "metadata": {},
   "source": [
    "Each POS and NER tag are represented by a number. In `dataset.features` we can see to which tag they refer (this [link](https://www.ling.upenn.edu/courses/Fall_2003/ling001/penn_treebank_pos.html) may serve you to look up the meaning)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "choice-artwork",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.features"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rational-jefferson",
   "metadata": {},
   "source": [
    "The following function will help us create the entities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "continent-silly",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_entities_POS(record):\n",
    "\n",
    "    entities = []\n",
    "    counter = 0\n",
    "\n",
    "    for i in range(len(record['pos_tags'])):\n",
    "\n",
    "        entity = (dataset.features[\"pos_tags\"].feature.names[record[\"pos_tags\"][i]], counter, counter + len(record[\"tokens\"][i]))\n",
    "        entities.append(entity)\n",
    "    \n",
    "        counter += len(record[\"tokens\"][i]) + 1\n",
    "        \n",
    "    return entities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "amended-appointment",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "records = []\n",
    "\n",
    "for record in dataset:\n",
    "    entities = parse_entities_POS(record)\n",
    "    records.append(rb.TokenClassificationRecord(\n",
    "        text=\" \".join(record[\"tokens\"]),\n",
    "        tokens=record[\"tokens\"],\n",
    "        annotation=entities,\n",
    "        annotation_agent=\"https://huggingface.co/datasets/conll2003\",\n",
    "        metadata={\n",
    "            \"split\": \"train\"\n",
    "            },\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "derived-billy",
   "metadata": {},
   "outputs": [],
   "source": [
    "rb.log(records=records, name=\"conll2003\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "assigned-camping",
   "metadata": {},
   "source": [
    "And so it is done! We have logged data from 5 different type of experiments, which now can be visualized in Rubrix UI "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "daily-halloween",
   "metadata": {},
   "source": [
    "## 3. Exploring predictions\n",
    "\n",
    "In this third part of the tutorial we are going to focus on loading predictions and annotations into Rubrix and visualize them from the UI. \n",
    "\n",
    "Rubrix let us play with the data in many different ways: visualizing by predicted class, by annotated class, by split, selecting which ones were wrongly classified, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "recreational-hormone",
   "metadata": {},
   "source": [
    "### Agnews and zeroshot classification\n",
    "\n",
    "To explore some logged data on Rubrix UI, we are going to predict the topic of some news with a zero-shot classifier (that we don't need to train), and compare the predicted category with the ground truth. The dataset we are going to use in this part is [ag_news](https://huggingface.co/datasets/ag_news), with information of over 1 million articles written in English.\n",
    "\n",
    "First of all, as always, we are going to load the dataset from 🤗 Hub and visualize its content."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compound-caution",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset  \n",
    "\n",
    "dataset = load_dataset(\"ag_news\", split='test[0:100]') # 20% is over 1500 records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "improving-printing",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "advisory-works",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.features"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "trained-worship",
   "metadata": {},
   "source": [
    "This dataset has articles from four different classes, so we can define a category list, which may come in handy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "upset-inspector",
   "metadata": {},
   "outputs": [],
   "source": [
    "categories = ['World', 'Sports', 'Business', 'Sci/Tech']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "electronic-retention",
   "metadata": {},
   "source": [
    "Now, it's time to load our zero-shot classification model. We present to options:\n",
    "\n",
    "1. [DistilBart-MNLI](https://huggingface.co/valhalla/distilbart-mnli-12-1)\n",
    "2. [squeezebert-mnli](\"typeform/squeezebert-mnli\")\n",
    "\n",
    "With the first model, the obtained results are probably going to be better, but it is a larger model, which could take longer to use. We are going to stick with the first one, but feel free to change it, and even to compare them!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "desperate-alexandria",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "model = \"valhalla/distilbart-mnli-12-1\"\n",
    "\n",
    "pl = pipeline('zero-shot-classification', model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "recovered-hawaiian",
   "metadata": {},
   "source": [
    "Let's try to make a quick prediction and take a look."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "continent-guest",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl(dataset[0]['text'], ['World', 'Sports', 'Business', 'Sci/Tech'], hypothesis_template='This example is {}.',multi_label=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "congressional-journal",
   "metadata": {},
   "source": [
    "Knowing how to make a prediction, we can now apply this to the whole selected dataset. Here, we also present you with two options:\n",
    "\n",
    "1. Traverse through all records in the dataset, predict each record and log it to Rubrix.\n",
    "2. Apply a map function to make the predictions and add that field to each record, and then log it as a whole to Rubrix.\n",
    "\n",
    "In the following categories, each approach is presented. You choose what you like the most, or even both (be careful with the time and the duplicated records, though!)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "criminal-correspondence",
   "metadata": {},
   "source": [
    "### First approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "discrete-briefing",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "for record in tqdm(dataset):\n",
    "    \n",
    "    # Make the prediction\n",
    "    model_output = pl(record['text'], categories, hypothesis_template='This example is {}.')\n",
    "\n",
    "    item = rb.TextClassificationRecord(\n",
    "        inputs={\"text\": record[\"text\"]},\n",
    "        prediction=list(zip(model_output['labels'], model_output['scores'])), \n",
    "        prediction_agent=\"https://huggingface.co/valhalla/distilbart-mnli-12-1\",\n",
    "        annotation=categories[record[\"label\"]],\n",
    "        annotation_agent=\"https://huggingface.co/datasets/ag_news\",\n",
    "        multi_label=True,\n",
    "        metadata={\n",
    "            \"split\": \"train\"\n",
    "            },\n",
    "        )\n",
    "    \n",
    "    \n",
    "    # Log to rubrix\n",
    "    rb.log(records=item, name=\"ag_news\")\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "immune-notebook",
   "metadata": {},
   "source": [
    "### Second approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "global-snowboard",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_predictions(records):\n",
    "    \n",
    "    predictions = pl([record for record in records['text']], categories, hypothesis_template='This example is {}.')\n",
    "    \n",
    "    \n",
    "    if isinstance(predictions, list):\n",
    "        return {\"labels_predicted\": [pred[\"labels\"] for pred in predictions], \"probabilities_predicted\": [pred[\"scores\"] for pred in predictions]}\n",
    "    else:\n",
    "        return {\"labels_predicted\": predictions[\"labels\"], \"probabilities_predicted\": predictions[\"scores\"]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "solar-minute",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_predicted = dataset.map(add_predictions, batched=True, batch_size=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "returning-collection",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_predicted[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compressed-newcastle",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "for record in tqdm(dataset_predicted):\n",
    "\n",
    "    item = rb.TextClassificationRecord(\n",
    "        inputs={\"text\": record[\"text\"]},\n",
    "        prediction=list(zip(record['labels_predicted'], record['probabilities_predicted'])), \n",
    "        prediction_agent=\"https://huggingface.co/valhalla/distilbart-mnli-12-1\",\n",
    "        annotation=categories[record[\"label\"]],\n",
    "        annotation_agent=\"https://huggingface.co/datasets/ag_news\",\n",
    "        multi_label=True,\n",
    "        metadata={\n",
    "            \"split\": \"train\"\n",
    "            },\n",
    "        )\n",
    "    \n",
    "    # Log to rubrix\n",
    "    rb.log(records=item, name=\"ag_news\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "936b0722",
   "metadata": {},
   "source": [
    "## Summary\n",
    "In this tutorial, we have learnt:\n",
    "\n",
    "- To log and explore NLP training datasets with the 🤗 `datasets` library.\n",
    "\n",
    "- To explore NLP predictions using a `zeroshot classifier` from the 🤗 `model hub`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69f98c1c-f8c7-4b08-9899-7d4c7be46d7a",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "### 📚 [Rubrix documentation](https://docs.rubrix.ml) for more guides and tutorials. \n",
    "\n",
    "### 🙋‍♀️ Join the Rubrix community! A good place to start is the [discussion forum](https://github.com/recognai/rubrix/discussions).\n",
    "\n",
    "### ⭐ Rubrix [Github repo](https://github.com/recognai/rubrix) to stay updated. "
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "b709380ea7d1cb2eb4650c0f11ac7e002ec6a534602815725771481b4784238c"
  },
  "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.5"
  },
  "metadata": {
   "interpreter": {
    "hash": "8c8e61cbe1071cbf939cab7691a0f9461f9f60562e343a6ebf81fec727c98593"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}