docs/guides/task_examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tasks Templates\n",
    "\n",
    "Hi there! In this article we wanted to share some examples of our supported tasks, so you can go from zero to hero as fast as possible. We are going to cover those tasks present in our [supported tasks list](../getting_started/supported_tasks.rst), so don't forget to stop by and take a look.\n",
    "\n",
    "The tasks are divided into their different category, from text classification to token classification. We will update this article, as well as the supported task list when a new task gets added to Rubrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text 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 labelled examples and seeing how they start predicting over the very same labels."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Text Categorization\n",
    "\n",
    "This is a general example of the Text Classification family of tasks. Here, we will try to assign pre-defined categories to sentences and texts. The possibilities are endless! Topic categorization, spam detection, and a vast etcétera.\n",
    "\n",
    "For our example, we are using the [SequeezeBERT](https://huggingface.co/typeform/squeezebert-mnli) zero-shot classifier for predicting the topic of a given text, in three different labels: politics, sports and technology. We are also using [AG](https://huggingface.co/datasets/ag_news), a collection of news, as our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "from transformers import pipeline\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Loading our dataset\n",
    "dataset = load_dataset(\"ag_news\", split=\"train[0:20]\")\n",
    "\n",
    "# Define our HuggingFace Pipeline\n",
    "classifier = pipeline(\n",
    "    \"zero-shot-classification\",\n",
    "    model=\"typeform/squeezebert-mnli\",\n",
    "    framework=\"pt\",\n",
    ")\n",
    "\n",
    "records = []\n",
    "\n",
    "for record in dataset:\n",
    "\n",
    "    # Making the prediction\n",
    "    prediction = classifier(\n",
    "        record[\"text\"],\n",
    "        candidate_labels=[\n",
    "            \"politics\",\n",
    "            \"sports\",\n",
    "            \"technology\",\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    # Creating the prediction entity as a list of tuples (label, probability)\n",
    "    prediction = list(zip(prediction[\"labels\"], prediction[\"scores\"]))\n",
    "\n",
    "    # Appending to the record list\n",
    "    records.append(\n",
    "        rb.TextClassificationRecord(\n",
    "            text=record[\"text\"],\n",
    "            prediction=prediction,\n",
    "            prediction_agent=\"https://huggingface.co/typeform/squeezebert-mnli\",\n",
    "            metadata={\"split\": \"train\"},\n",
    "        )\n",
    "    )\n",
    "\n",
    "# Logging into Rubrix\n",
    "rb.log(\n",
    "    records=records,\n",
    "    name=\"text-categorization\",\n",
    "    tags={\n",
    "        \"task\": \"text-categorization\",\n",
    "        \"phase\": \"data-analysis\",\n",
    "        \"family\": \"text-classification\",\n",
    "        \"dataset\": \"ag_news\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sentiment Analysis\n",
    "\n",
    "In this kind of project, we want our models to be able to detect the polarity of the input. Categories like *positive*, *negative* or *neutral* are often used. \n",
    "\n",
    "For this example, we are going to use an [Amazon review polarity dataset](https://huggingface.co/datasets/amazon_polarity), and a sentiment analysis [roBERTa model](https://huggingface.co/cardiffnlp/twitter-roberta-base-sentiment?text=I+like+you.+I+love+you), which returns `LABEL 0` for positive, `LABEL 1` for neutral and `LABEL 2` for negative. We will handle that in the code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "from transformers import pipeline\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Loading our dataset\n",
    "dataset = load_dataset(\"amazon_polarity\", split=\"train[0:20]\")\n",
    "\n",
    "# Define our HuggingFace Pipeline\n",
    "classifier = pipeline(\n",
    "    \"text-classification\",\n",
    "    model=\"cardiffnlp/twitter-roberta-base-sentiment\",\n",
    "    framework=\"pt\",\n",
    "    return_all_scores=True,\n",
    ")\n",
    "\n",
    "# Make a dictionary to translate labels to a friendly-language\n",
    "translate_labels = {\n",
    "    \"LABEL_0\": \"positive\",\n",
    "    \"LABEL_1\": \"neutral\",\n",
    "    \"LABEL_2\": \"negative\",\n",
    "}\n",
    "\n",
    "records = []\n",
    "\n",
    "for record in dataset:\n",
    "\n",
    "    # Making the prediction\n",
    "    predictions = classifier(\n",
    "        record[\"content\"],\n",
    "    )\n",
    "\n",
    "    # Creating the prediction entity as a list of tuples (label, probability)\n",
    "    prediction = [\n",
    "        (translate_labels[prediction[\"label\"]], prediction[\"score\"])\n",
    "        for prediction in predictions[0]\n",
    "    ]\n",
    "\n",
    "    # Appending to the record list\n",
    "    records.append(\n",
    "        rb.TextClassificationRecord(\n",
    "            text=record[\"content\"],\n",
    "            prediction=prediction,\n",
    "            prediction_agent=\"https://huggingface.co/cardiffnlp/twitter-roberta-base-sentiment\",\n",
    "            metadata={\"split\": \"train\"},\n",
    "        )\n",
    "    )\n",
    "\n",
    "# Logging into Rubrix\n",
    "rb.log(\n",
    "    records=records,\n",
    "    name=\"sentiment-analysis\",\n",
    "    tags={\n",
    "        \"task\": \"sentiment-analysis\",\n",
    "        \"phase\": \"data-annotation\",\n",
    "        \"family\": \"text-classification\",\n",
    "        \"dataset\": \"amazon-polarity\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Semantic Textual Similarity\n",
    "\n",
    "This task is all about how close or far a given text is from any other. We want models that output a value of closeness between two inputs.\n",
    "\n",
    "For our example, we will be using [MRPC dataset](https://paperswithcode.com/dataset/mrpc), a corpus consisting of 5,801 sentence pairs collected from newswire articles. These pairs could (or could not) be paraphrases. Our model will be a [sentence Transformer](https://huggingface.co/sentence-transformers/paraphrase-MiniLM-L12-v2), trained specifically for this task. \n",
    "\n",
    "As HuggingFace Transformers does not support natively this task, we will be using the [Sentence Transformer](https://www.sbert.net) framework. For more information about how to make these predictions with HuggingFace Transformer, please visit this [link](https://huggingface.co/sentence-transformers/paraphrase-MiniLM-L12-v2)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "from sentence_transformers import SentenceTransformer, util\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Loading our dataset\n",
    "dataset = load_dataset(\"glue\", \"mrpc\", split=\"train[0:20]\")\n",
    "\n",
    "# Loading the model\n",
    "model = SentenceTransformer(\"paraphrase-MiniLM-L6-v2\")\n",
    "\n",
    "records = []\n",
    "\n",
    "for record in dataset:\n",
    "\n",
    "    # Creating a sentence list\n",
    "    sentences = [record[\"sentence1\"], record[\"sentence2\"]]\n",
    "\n",
    "    # Obtaining similarity\n",
    "    paraphrases = util.paraphrase_mining(model, sentences)\n",
    "\n",
    "    for paraphrase in paraphrases:\n",
    "        score, _, _ = paraphrase\n",
    "\n",
    "    # Building up the prediction tuples\n",
    "    prediction = [(\"similar\", score), (\"not similar\", 1 - score)]\n",
    "\n",
    "    # Appending to the record list\n",
    "    records.append(\n",
    "        rb.TextClassificationRecord(\n",
    "            inputs={\n",
    "                \"sentence 1\": record[\"sentence1\"],\n",
    "                \"sentence 2\": record[\"sentence2\"],\n",
    "            },\n",
    "            prediction=prediction,\n",
    "            prediction_agent=\"https://huggingface.co/sentence-transformers/paraphrase-MiniLM-L12-v2\",\n",
    "            metadata={\"split\": \"train\"},\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "# Logging into Rubrix\n",
    "rb.log(\n",
    "    records=records,\n",
    "    name=\"semantic-textual-similarity\",\n",
    "    tags={\n",
    "        \"task\": \"similarity\",\n",
    "        \"type\": \"paraphrasing\",\n",
    "        \"family\": \"text-classification\",\n",
    "        \"dataset\": \"mrpc\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Natural Language Inference\n",
    "\n",
    "Natural language inference is the task of determining whether a hypothesis is true (which will mean entailment), false (contradiction), or undetermined (neutral) given a premise. This task also works with pair of sentences. \n",
    "\n",
    "Our dataset will be the famous [SNLI](https://huggingface.co/datasets/snli), a collection of 570k human-written English sentence pairs; and our model will be a [zero-shot, cross encoder for inference](https://huggingface.co/cross-encoder/nli-MiniLM2-L6-H768)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "from transformers import pipeline\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Loading our dataset\n",
    "dataset = load_dataset(\"snli\", split=\"train[0:20]\")\n",
    "\n",
    "# Define our HuggingFace Pipeline\n",
    "classifier = pipeline(\n",
    "    \"zero-shot-classification\",\n",
    "    model=\"cross-encoder/nli-MiniLM2-L6-H768\",\n",
    "    framework=\"pt\",\n",
    ")\n",
    "\n",
    "records = []\n",
    "\n",
    "for record in dataset:\n",
    "\n",
    "    # Making the prediction\n",
    "    prediction = classifier(\n",
    "        record[\"premise\"] + record[\"hypothesis\"],\n",
    "        candidate_labels=[\n",
    "            \"entailment\",\n",
    "            \"contradiction\",\n",
    "            \"neutral\",\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    # Creating the prediction entity as a list of tuples (label, probability)\n",
    "    prediction = list(zip(prediction[\"labels\"], prediction[\"scores\"]))\n",
    "\n",
    "    # Appending to the record list\n",
    "    records.append(\n",
    "        rb.TextClassificationRecord(\n",
    "            inputs={\"premise\": record[\"premise\"], \"hypothesis\": record[\"hypothesis\"]},\n",
    "            prediction=prediction,\n",
    "            prediction_agent=\"https://huggingface.co/cross-encoder/nli-MiniLM2-L6-H768\",\n",
    "            metadata={\"split\": \"train\"},\n",
    "        )\n",
    "    )\n",
    "\n",
    "# Logging into Rubrix\n",
    "rb.log(\n",
    "    records=records,\n",
    "    name=\"natural-language-inference\",\n",
    "    tags={\n",
    "        \"task\": \"nli\",\n",
    "        \"family\": \"text-classification\",\n",
    "        \"dataset\": \"snli\",\n",
    "    },\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stance Detection\n",
    "\n",
    "Stance detection is the NLP task which seeks to extract from a subject's reaction to a claim made by a primary actor. It is a core part of a set of approaches to fake news assessment. For example:\n",
    "\n",
    "-   **Source**: \"*Apples are the most delicious fruit in existence*\"\n",
    "-   **Reply**: \"*Obviously not, because that is a reuben from Katz's*\"\n",
    "-   **Stance**: deny\n",
    "\n",
    "But it can be done in many different ways. In the search of fake news, there is usually one source of text.\n",
    "\n",
    "We will be using the [LIAR datastet](https://huggingface.co/datasets/liar), a fake news detection dataset with 12.8K human labeled short statements from politifact.com's API, and each statement is evaluated by a politifact.com editor for its truthfulness, and a zero-shot [distilbart](https://huggingface.co/valhalla/distilbart-mnli-12-3) model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "from transformers import pipeline\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Loading our dataset\n",
    "dataset = load_dataset(\"liar\", split=\"train[0:20]\")\n",
    "\n",
    "# Define our HuggingFace Pipeline\n",
    "classifier = pipeline(\n",
    "    \"zero-shot-classification\",\n",
    "    model=\"valhalla/distilbart-mnli-12-3\",\n",
    "    framework=\"pt\",\n",
    ")\n",
    "\n",
    "records = []\n",
    "\n",
    "for record in dataset:\n",
    "\n",
    "    # Making the prediction\n",
    "    prediction = classifier(\n",
    "        record[\"statement\"],\n",
    "        candidate_labels=[\n",
    "            \"false\",\n",
    "            \"half-true\",\n",
    "            \"mostly-true\",\n",
    "            \"true\",\n",
    "            \"barely-true\",\n",
    "            \"pants-fire\",\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    # Creating the prediction entity as a list of tuples (label, probability)\n",
    "    prediction = list(zip(prediction[\"labels\"], prediction[\"scores\"]))\n",
    "\n",
    "    # Appending to the record list\n",
    "    records.append(\n",
    "        rb.TextClassificationRecord(\n",
    "            text=record[\"statement\"],\n",
    "            prediction=prediction,\n",
    "            prediction_agent=\"https://huggingface.co/typeform/squeezebert-mnli\",\n",
    "            metadata={\"split\": \"train\"},\n",
    "        )\n",
    "    )\n",
    "\n",
    "# Logging into Rubrix\n",
    "rb.log(\n",
    "    records=records,\n",
    "    name=\"stance-detection\",\n",
    "    tags={\n",
    "        \"task\": \"stance detection\",\n",
    "        \"family\": \"text-classification\",\n",
    "        \"dataset\": \"liar\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multilabel Text Classification\n",
    "\n",
    "A variation of the text classification basic problem, in this task we want to categorize a given input into one or more categories. The labels or categories are not mutually exclusive.\n",
    "\n",
    "For this example, we will be using the [go emotions](https://huggingface.co/datasets/go_emotions) dataset, with Reddit comments categorized in 27 different emotions. Alongside the dataset, we've chosen a [DistilBERT model](https://huggingface.co/joeddav/distilbert-base-uncased-go-emotions-student), distilled from a zero-shot classification pipeline.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "from transformers import pipeline\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Loading our dataset\n",
    "dataset = load_dataset(\"go_emotions\", split=\"train[0:20]\")\n",
    "\n",
    "# Define our HuggingFace Pipeline\n",
    "classifier = pipeline(\n",
    "    \"text-classification\",\n",
    "    model=\"joeddav/distilbert-base-uncased-go-emotions-student\",\n",
    "    framework=\"pt\",\n",
    "    return_all_scores=True,\n",
    ")\n",
    "\n",
    "records = []\n",
    "\n",
    "for record in dataset:\n",
    "\n",
    "    # Making the prediction\n",
    "    prediction = classifier(record[\"text\"], multi_label=True)\n",
    "\n",
    "    # Creating the prediction entity as a list of tuples (label, probability)\n",
    "    prediction = [(pred[\"label\"], pred[\"score\"]) for pred in prediction[0]]\n",
    "\n",
    "    # Appending to the record list\n",
    "    records.append(\n",
    "        rb.TextClassificationRecord(\n",
    "            text=record[\"text\"],\n",
    "            prediction=prediction,\n",
    "            prediction_agent=\"https://huggingface.co/typeform/squeezebert-mnli\",\n",
    "            metadata={\"split\": \"train\"},\n",
    "            multi_label=True,  # we also need to set the multi_label option in Rubrix\n",
    "        )\n",
    "    )\n",
    "\n",
    "# Logging into Rubrix\n",
    "rb.log(\n",
    "    records=records,\n",
    "    name=\"multilabel-text-classification\",\n",
    "    tags={\n",
    "        \"task\": \"multilabel-text-classification\",\n",
    "        \"family\": \"text-classification\",\n",
    "        \"dataset\": \"go_emotions\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Node Classification\n",
    "\n",
    "The node classification task is the one where the model has to determine the labelling of samples (represented as nodes) by looking at the labels of their neighbours, in a Graph Neural Network. If you want to know more about GNNs, we've made a [tutorial](https://docs.rubrix.ml/en/stable/tutorials/03-kglab_pytorch_geometric.html) about them using Kglab and PyTorch Geometric, which integrates Rubrix into the pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Token Classification\n",
    "\n",
    "Token classification kind-of-tasks are NLP tasks aimed to divide the input text into words, or syllables, and assign certain values to them. Think about giving each word in a sentence its grammatical category, or highlight which parts of a medical report belong to a certain speciality. There are some popular ones like NER or POS-tagging. For this part of the article, we will use [spaCy](https://spacy.io/) with Rubrix to track and monitor Token Classification tasks.\n",
    "\n",
    "Remember to install spaCy and datasets, or running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install datasets -qqq\n",
    "%pip install -U spacy -qqq\n",
    "%pip install protobuf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NER\n",
    "Named entity recognition (NER) is the task of tagging entities in text with their corresponding type. Approaches typically use *BIO* notation, which differentiates the beginning (**B**) and the inside (**I**) of entities. **O** is used for non-entity tokens.\n",
    "\n",
    "For this tutorial, we're going to use the [*Gutenberg Time*](https://huggingface.co/datasets/gutenberg_time) dataset from the Hugging Face Hub. It contains all explicit time references in a dataset of 52,183 novels whose full text is available via Project Gutenberg. From extracts of novels, we are surely going to find some NER entities. We will also use the `en_core_web_trf` pretrained English model, a Roberta-based spaCy model. If you do not have them installed, run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy download en_core_web_trf #Download the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "import spacy\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Load our dataset\n",
    "dataset = load_dataset(\"gutenberg_time\", split=\"train[0:20]\")\n",
    "\n",
    "# Load the spaCy model\n",
    "nlp = spacy.load(\"en_core_web_trf\")\n",
    "\n",
    "records = []\n",
    "\n",
    "for record in dataset:\n",
    "\n",
    "    # We only need the text of each instance\n",
    "    text = record[\"tok_context\"]\n",
    "\n",
    "    # spaCy Doc creation\n",
    "    doc = nlp(text)\n",
    "\n",
    "    # Prediction entities with the tuples (label, start character, end character)\n",
    "    entities = [(ent.label_, ent.start_char, ent.end_char) for ent in doc.ents]\n",
    "\n",
    "    # Pre-tokenized input text\n",
    "    tokens = [token.text for token in doc]\n",
    "\n",
    "    # Rubrix TokenClassificationRecord list\n",
    "    records.append(\n",
    "        rb.TokenClassificationRecord(\n",
    "            text=text,\n",
    "            tokens=tokens,\n",
    "            prediction=entities,\n",
    "            prediction_agent=\"en_core_web_trf\",\n",
    "        )\n",
    "    )\n",
    "\n",
    "# Logging into Rubrix\n",
    "rb.log(\n",
    "    records=records,\n",
    "    name=\"ner\",\n",
    "    tags={\n",
    "        \"task\": \"NER\",\n",
    "        \"family\": \"token-classification\",\n",
    "        \"dataset\": \"gutenberg-time\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### POS tagging\n",
    "\n",
    "A POS tag (or part-of-speech tag) is a special label assigned to each word in a text corpus to indicate the part of speech and often also other grammatical categories such as tense, number, case etc. POS tags are used in corpus searches and in-text analysis tools and algorithms.\n",
    "\n",
    "We will be repeating duo for this second spaCy example, with the [*Gutenberg Time*](https://huggingface.co/datasets/gutenberg_time) dataset from the Hugging Face Hub and the `en_core_web_trf` pretrained English model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "import spacy\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Load our dataset\n",
    "dataset = load_dataset(\"gutenberg_time\", split=\"train[0:10]\")\n",
    "\n",
    "# Load the spaCy model\n",
    "nlp = spacy.load(\"en_core_web_trf\")\n",
    "\n",
    "records = []\n",
    "\n",
    "for record in dataset:\n",
    "\n",
    "    # We only need the text of each instance\n",
    "    text = record[\"tok_context\"]\n",
    "\n",
    "    # spaCy Doc creation\n",
    "    doc = nlp(text)\n",
    "\n",
    "    # Creating the prediction entity as a list of tuples (tag, start_char, end_char)\n",
    "    prediction = [(token.pos_, token.idx, token.idx + len(token)) for token in doc]\n",
    "\n",
    "    # Rubrix TokenClassificationRecord list\n",
    "    records.append(\n",
    "        rb.TokenClassificationRecord(\n",
    "            text=text,\n",
    "            tokens=[token.text for token in doc],\n",
    "            prediction=prediction,\n",
    "            prediction_agent=\"en_core_web_trf\",\n",
    "        )\n",
    "    )\n",
    "\n",
    "# Logging into Rubrix\n",
    "rb.log(\n",
    "    records=records,\n",
    "    name=\"pos-tagging\",\n",
    "    tags={\n",
    "        \"task\": \"pos-tagging\",\n",
    "        \"family\": \"token-classification\",\n",
    "        \"dataset\": \"gutenberg-time\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Slot Filling\n",
    "\n",
    "The goal of Slot Filling is to identify, from a running dialog different slots, which one correspond to different parameters of the user’s query. For instance, when a user queries for nearby restaurants, key slots for location and preferred food are required for a dialog system to retrieve the appropriate information. Thus, the goal is to look for specific pieces of information in the request and tag the corresponding tokens accordingly.\n",
    "\n",
    "We made a tutorial on this matter for our open-source NLP library, [biome.text](https://recognai.github.io/biome-text/v3.0.0/). We will use similar procedures here, focusing on the logging of the information. If you want to see in-depth explanations on how the pipelines are made, please visit [the tutorial](https://recognai.github.io/biome-text/v3.0.0/documentation/tutorials/2-Training_a_sequence_tagger_for_Slot_Filling.html#training-a-sequence-tagger-for-slot-filling)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by downloading biome.text and importing it alongside Rubrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U biome-text\n",
    "exit(0)  # Force restart of the runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "\n",
    "from biome.text import Pipeline, Dataset, PipelineConfiguration, VocabularyConfiguration, Trainer\n",
    "from biome.text.configuration import FeaturesConfiguration, WordFeatures, CharFeatures\n",
    "from biome.text.modules.configuration import Seq2SeqEncoderConfiguration\n",
    "from biome.text.modules.heads import TokenClassificationConfiguration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this tutorial we will use the [SNIPS data set](https://github.com/sonos/nlu-benchmark/tree/master/2017-06-custom-intent-engines) adapted by [Su Zhu](https://github.com/sz128/slot_filling_and_intent_detection_of_SLU/tree/master/data/snips)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!curl -O https://biome-tutorials-data.s3-eu-west-1.amazonaws.com/token_classifier/train.json\n",
    "!curl -O https://biome-tutorials-data.s3-eu-west-1.amazonaws.com/token_classifier/valid.json\n",
    "!curl -O https://biome-tutorials-data.s3-eu-west-1.amazonaws.com/token_classifier/test.json\n",
    "\n",
    "train_ds = Dataset.from_json(\"train.json\")\n",
    "valid_ds = Dataset.from_json(\"valid.json\")\n",
    "test_ds = Dataset.from_json(\"test.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Afterwards, we need to configure our biome.text Pipeline. More information on this configuration [here](https://recognai.github.io/biome-text/v3.0.0/documentation/tutorials/2-Training_a_sequence_tagger_for_Slot_Filling.html#configure-your-biome-text-pipeline)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_feature = WordFeatures(\n",
    "    embedding_dim=300,\n",
    "    weights_file=\"https://dl.fbaipublicfiles.com/fasttext/vectors-english/wiki-news-300d-1M.vec.zip\",\n",
    ")\n",
    "\n",
    "char_feature = CharFeatures(\n",
    "    embedding_dim=32,\n",
    "    encoder={\n",
    "        \"type\": \"gru\",\n",
    "        \"bidirectional\": True,\n",
    "        \"num_layers\": 1,\n",
    "        \"hidden_size\": 32,\n",
    "    },\n",
    "    dropout=0.1\n",
    ")\n",
    "\n",
    "features_config = FeaturesConfiguration(\n",
    "    word=word_feature, \n",
    "    char=char_feature\n",
    ")\n",
    "\n",
    "encoder_config = Seq2SeqEncoderConfiguration(\n",
    "    type=\"gru\",\n",
    "    bidirectional=True,\n",
    "    num_layers=1,\n",
    "    hidden_size=128,\n",
    ")\n",
    "\n",
    "labels = {tag[2:] for tags in train_ds[\"labels\"] for tag in tags if tag != \"O\"}\n",
    "\n",
    "for ds in [train_ds, valid_ds, test_ds]:\n",
    "    ds.rename_column_(\"labels\", \"tags\")\n",
    "\n",
    "head_config = TokenClassificationConfiguration(\n",
    "    labels=list(labels),\n",
    "    label_encoding=\"BIO\",\n",
    "    top_k=1,\n",
    "    feedforward={\n",
    "        \"num_layers\": 1,\n",
    "        \"hidden_dims\": [128],\n",
    "        \"activations\": [\"relu\"],\n",
    "        \"dropout\": [0.1],\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now, let's train our model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline_config = PipelineConfiguration(\n",
    "    name=\"slot_filling_tutorial\",\n",
    "    features=features_config,\n",
    "    encoder=encoder_config,\n",
    "    head=head_config,\n",
    ")\n",
    "\n",
    "pl = Pipeline.from_config(pipeline_config)\n",
    "\n",
    "vocab_config = VocabularyConfiguration(min_count={\"word\": 2}, include_valid_data=True)\n",
    "\n",
    "trainer = Trainer(\n",
    "    pipeline=pl,\n",
    "    train_dataset=train_ds,\n",
    "    valid_dataset=valid_ds,\n",
    "    vocab_config=vocab_config,\n",
    "    trainer_config=None,\n",
    ")\n",
    "\n",
    "trainer.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having trained our model, we can go ahead and log the predictions to Rubrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = Dataset.from_json(\"test.json\")\n",
    "\n",
    "records = []\n",
    "\n",
    "for record in dataset[0:10][\"text\"]:\n",
    "\n",
    "    # We only need the text of each instance\n",
    "    text = \" \".join(word for word in record)\n",
    "\n",
    "    # Predicting tags and entities given the input text\n",
    "    prediction = pl.predict(text=text)\n",
    "\n",
    "    # Creating the prediction entity as a list of tuples (tag, start_char, end_char)\n",
    "    prediction = [\n",
    "        (token[\"label\"], token[\"start\"], token[\"end\"])\n",
    "        for token in prediction[\"entities\"][0]\n",
    "    ]\n",
    "\n",
    "    # Rubrix TokenClassificationRecord list\n",
    "    records.append(\n",
    "        rb.TokenClassificationRecord(\n",
    "            text=text,\n",
    "            tokens=record,\n",
    "            prediction=prediction,\n",
    "            prediction_agent=\"biome_slot_filling_tutorial\",\n",
    "        )\n",
    "    )\n",
    "\n",
    "# Logging into Rubrix\n",
    "rb.log(\n",
    "    records=records,\n",
    "    name=\"slot-filling\",\n",
    "    tags={\n",
    "        \"task\": \"slot-filling\",\n",
    "        \"family\": \"token-classification\",\n",
    "        \"dataset\": \"SNIPS\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Text2Text\n",
    "\n",
    "The expression *Text2Text* encompasses text generation tasks where the model receives and outputs a sequence of tokens.\n",
    "Examples of such tasks are machine translation, text summarization, paraphrase generation, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Machine translation\n",
    "\n",
    "Machine translation is the task of translating text from one language to another.\n",
    "It is arguably one of the oldest NLP tasks, but human parity remains an [open challenge](https://aclanthology.org/W18-6312.pdf) especially for low resource languages and domains.\n",
    "\n",
    "In the following small example we will showcase how *Rubrix* can help you to fine-tune an English-to-Spanish translation model.\n",
    "Let us assume we want to translate \"Sesame Street\" related content.\n",
    "If you have been to Spain before you probably noticed that named entities (like character or band names) are often translated quite literally or are very different from the original ones.  \n",
    "We will use a pre-trained 🤗 transformers model to get a few suggestions for the translation, and then correct them in *Rubrix* to obtain a training set for the fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#!pip install transformers\n",
    "\n",
    "from transformers import pipeline\n",
    "import rubrix as rb\n",
    "\n",
    "# Instantiate the translator\n",
    "translator = pipeline(\"translation_en_to_es\", model=\"Helsinki-NLP/opus-mt-en-es\")\n",
    "\n",
    "# 'Sesame Street' related phrase\n",
    "en_phrase = \"Sesame Street is an American educational children's television series starring the muppets Ernie and Bert.\"\n",
    "\n",
    "# Get two predictions from the translator\n",
    "es_predictions = [output[\"translation_text\"] for output in translator(en_phrase, num_return_sequences=2)]\n",
    "\n",
    "# Log the record to Rubrix and correct them\n",
    "record = rb.Text2TextRecord(\n",
    "    text=en_phrase,\n",
    "    prediction=es_predictions,\n",
    ")\n",
    "rb.log(record, name=\"sesame_street_en-es\")\n",
    "\n",
    "# For a real training set you probably would need more than just one 'Sesame Street' related phrase."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the *Rubrix* web app we can now easily browse the predictions and annotate the records with a corrected prediction of our choice.\n",
    "The predictions for our example phrase are:\n",
    "1. Sesame Street es una serie de televisión infantil estadounidense protagonizada por los muppets Ernie y Bert.\n",
    "2. Sesame Street es una serie de televisión infantil y educativa estadounidense protagonizada por los muppets Ernie y Bert.\n",
    "\n",
    "We probably would choose the second one and correct it in the following way:\n",
    "\n",
    "2. *Barrio Sésamo* es una serie de televisión infantil y educativa estadounidense protagonizada por los *teleñecos Epi y Blas*.*\n",
    "\n",
    "\n",
    "After correcting a substantial number of example phrases, we can load the corrected data set as a DataFrame to use it for the fine-tuning of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load corrected translations for the fine-tuning of the translation model\n",
    "df = rb.load(\"sesame_street_en-es\")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "4cac2ad44382dcbde9c9d45667b9ac0fec163e57feefe7fa8ea5d11fd16eb612"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}