notebooks/W2T_Intro_Transformers_Datasets.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/gsarti/ik-nlp-tutorials/blob/main/notebooks/W2T_Intro_Transformers_Datasets.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run in Colab to install local packages\n",
    "!pip install transformers sentencepiece torch datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to 🤗 Transformers and 🤗 Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*This tutorial is based off some chapters of the [HuggingFace Course](https://huggingface.co/course/chapter1/1), take a look for a more detailed overview!*\n",
    "\n",
    "Transformer models are nowadays the state-of-the-art and de-facto standard to solve all kinds of NLP tasks, from tagging to machine translation, to text classification.\n",
    "\n",
    "The usage of these models has been widely simplified and democratized by [HuggingFace](https://huggingface.co/) (🤗 in short), the startup behind the popular [🤗 Transformers library](https://huggingface.co/transformers/). The 🤗 Transformers library is completely open-source and provides a unified framework to create, train and use many transformer-based models, accompanied by a Cloud-hosting service called [Model hub](https://huggingface.co/models) (similar to Pytorch and Tensorflow Hubs) in which every user can host and share pre-trained and fine-tuned open-source models, and which currently contains over 25k models (as of Jan 24, 2022).\n",
    "\n",
    "We are going to start with a quick overview of the 🤗 Transformers library and its usage, and then we will dive into the 🤗 Datasets library, which is the largest open collection of text datasets ready for usage with 🤗 Transformers and other machine learning frameworks (also hosted on the [Dataset hub](https://huggingface.co/datasets)) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pipelines\n",
    "\n",
    "The most basic object in the 🤗 Transformers library is the `pipeline` function. It connects a model with its necessary preprocessing and postprocessing steps, allowing us to directly input any text and get an intelligible answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to distilbert-base-uncased-finetuned-sst-2-english (https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'label': 'NEGATIVE', 'score': 0.9915691018104553}]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "classifier = pipeline(\"sentiment-analysis\")\n",
    "classifier(\"I've been waiting for the IK-NLP course for my whole life.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Multiple sentences can also be passed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'label': 'POSITIVE', 'score': 0.9807901382446289},\n",
       " {'label': 'NEGATIVE', 'score': 0.9996439218521118}]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classifier(\n",
    "    [\"I've been waiting for this course for my whole life.\", \"I hate this course so much!\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, this pipeline selects a particular pretrained model that has been fine-tuned for sentiment analysis in English (`distilbert-base-uncased-finetuned-sst-2-english`). The model is downloaded and cached when you create the classifier object. If you rerun the command, the cached model will be used instead and there is no need to download the model again.\n",
    "\n",
    "There are three main steps involved when you pass some text to a pipeline:\n",
    "\n",
    "- The text is preprocessed into a format the model can understand.\n",
    "- The preprocessed inputs are passed to the model.\n",
    "- The predictions of the model are post-processed, so you can make sense of them.\n",
    "\n",
    "Some of the [currently available pipelines](https://huggingface.co/transformers/main_classes/pipelines.html) are:\n",
    "\n",
    "- `sentiment-analysis`\n",
    "- `text-generation`\n",
    "- `fill-mask` (filling a masked token or span with a predicted one)\n",
    "- `text2text-generation`\n",
    "- `ner` (named entity recognition)\n",
    "- `question-answering`\n",
    "\n",
    "Some pipelines, such as `sentiment-analysis`, `translation` and `summarization` are abstractions over other, more general ones (e.g. `summarization` and `translation` are abstractions over `text2text-generation`, `sentiment-analysis` over `text-classification`).\n",
    "\n",
    "Let's see some examples of pipelines in action."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Text generation \n",
    "\n",
    "The text generation (i.e. autoregressive language modeling) setting has become widely popularized by models such as [GPT-3](https://en.wikipedia.org/wiki/GPT-3). Given a prompt, the model will auto-complete it by generating the remaining text. This is similar to the predictive text feature that is found on many phones. Here is an example using a small [GPT-2](https://huggingface.co/gpt2) model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to gpt2 (https://huggingface.co/gpt2)\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'generated_text': 'The goal of the course is to ensure that students are familiar with a number of fundamental techniques and algorithms in the area of natural language processing, such as: vernacular translation for example\\n\\nsimplifiers for sentences\\n\\ninterlinear transformations\\n'}]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "generator = pipeline(\"text-generation\")\n",
    "generator(\n",
    "    \"The goal of the course is to ensure that students are familiar with a number of fundamental \"\n",
    "    \"techniques and algorithms in the area of natural language processing, such as: \"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ">🤡 **Fun Fact**: Language models trained to perform autoregressive language modeling are already widely used in the industry. For example, the [Github Copilot](https://copilot.github.com/) integrated to the VisualStudio Code editor is a model trained to autocomplete text and code snippets, and is currently helping me in writing these notebooks (grey text are model suggestions).\n",
    "\n",
    "![Github Copilot autocompletion](https://github.com/gsarti/ik-nlp-tutorials/raw/main/img/copilot.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try controlling how many different sequences are generated with the argument `num_return_sequences` and the total maximal length of the output text with the argument `max_length`.\n",
    "\n",
    "We can use any custom model from the [Model hub](https://huggingface.co/models) by simply passing it (or its identifier) when creating the pipeline. We'll now use a [Dutch GPT-2 model](https://huggingface.co/GroNLP/gpt2-small-dutch) pretrained by our colleagues to generate some Dutch text:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/gsarti/projects/venv/lib/python3.8/site-packages/transformers/generation_utils.py:2142: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  next_indices = next_tokens // vocab_size\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'generated_text': \"Het doel van de cursus is ervoor te zorgen dat studenten bekend zijn met een aantal fundamentele technieken en algoritmen op het gebied van natuurlijke taalverwerking, zoals:\\nBeslissingen die niet tot uitdrukking kunnen worden gebracht of door middel van andere methoden verkregen. In deze context wordt er meestal gebruikgemaakt van 'informatief' interpretaties (bijvoorbeeld bij zinnen uit verschillende talen). De belangrijkste manier om dergelijke interpretaties in praktijk te brengen is bijvoorbeeld als volgt:\\nAls je wilt zeggen wat iemand zegt, kan\"}]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "generator = pipeline(\"text-generation\", model=\"GroNLP/gpt2-small-dutch\")\n",
    "generator(\n",
    "    \"Het doel van de cursus is ervoor te zorgen dat studenten bekend zijn met een aantal \"\n",
    "    \"fundamentele technieken en algoritmen op het gebied van natuurlijke taalverwerking, zoals:\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can refine your search for a model by clicking on the language tags, and pick a model that will generate text in another language. The Model Hub even contains checkpoints for multilingual models that support several languages.\n",
    "\n",
    "Once you select a model by clicking on it, you’ll see that there is a widget enabling you to try it directly online. This way you can quickly test the model’s capabilities before downloading it. More info on text generation here: [https://huggingface.co/tasks/text-generation](https://huggingface.co/tasks/text-generation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mask-filling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `fill-mask` (i.e. masked language modeling) pipeline is used to fill masked tokens with a predicted token, which is a common pre-training task for encoder-only transformers leveraging bidirectional context such as [BERT](https://huggingface.co/bert-base-uncased). This is useful to fill gaps in the text with the most likely answer given the context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to distilroberta-base (https://huggingface.co/distilroberta-base)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'sequence': 'This course will teach you all about natural language processing models.',\n",
       "  'score': 0.5701010823249817,\n",
       "  'token': 1632,\n",
       "  'token_str': ' natural'},\n",
       " {'sequence': 'This course will teach you all about functional language processing models.',\n",
       "  'score': 0.06381034106016159,\n",
       "  'token': 12628,\n",
       "  'token_str': ' functional'},\n",
       " {'sequence': 'This course will teach you all about programming language processing models.',\n",
       "  'score': 0.043610505759716034,\n",
       "  'token': 8326,\n",
       "  'token_str': ' programming'}]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "unmasker = pipeline(\"fill-mask\")\n",
    "unmasker(\"This course will teach you all about <mask> language processing models.\", top_k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `top_k` argument controls how many possibilities you want to be displayed. Note that here the model fills in the special `<mask>` word, which is often referred to as a mask token. Other mask-filling models might have different mask tokens, so it’s always good to verify the proper mask word when exploring other models. One way to check it is by looking at the mask word used in the widget. More info on mask-filling here: [https://huggingface.co/tasks/fill-mask](https://huggingface.co/tasks/fill-mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Text2Text Generation\n",
    "\n",
    "The `text2text-generation` pipeline encompasses all the sequence-to-sequence tasks on which a model was trained. We're gonna cover Encoder-Decoder architectures in week 6, but in the meantime you can see here some examples of sequence-to-sequence tasks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'generated_text': ' America has changed dramatically during recent years . The number of engineering graduates in the U.S. has declined in traditional engineering disciplines such as mechanical, civil,    electrical, chemical, and aeronautical engineering . Rapidly developing economies such as China and India continue to encourage and advance the teaching of engineering .'}]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Summarization: reduce a text to its summary\n",
    "# The same result can be achieved using the `summarization` pipeline.\n",
    "from transformers import pipeline\n",
    "\n",
    "summarizer = pipeline(\"text2text-generation\", model=\"sshleifer/distilbart-cnn-12-6\")\n",
    "summarizer(\n",
    "    \"\"\"\n",
    "    America has changed dramatically during recent years. Not only has the number of \n",
    "    graduates in traditional engineering disciplines such as mechanical, civil, \n",
    "    electrical, chemical, and aeronautical engineering declined, but in most of \n",
    "    the premier American universities engineering curricula now concentrate on \n",
    "    and encourage largely the study of engineering science. As a result, there \n",
    "    are declining offerings in engineering subjects dealing with infrastructure, \n",
    "    the environment, and related issues, and greater concentration on high \n",
    "    technology subjects, largely supporting increasingly complex scientific \n",
    "    developments. While the latter is important, it should not be at the expense \n",
    "    of more traditional engineering.\n",
    "\n",
    "    Rapidly developing economies such as China and India, as well as other \n",
    "    industrial countries in Europe and Asia, continue to encourage and advance \n",
    "    the teaching of engineering. Both China and India, respectively, graduate \n",
    "    six and eight times as many traditional engineers as does the United States. \n",
    "    Other industrial countries at minimum maintain their output, while America \n",
    "    suffers an increasingly serious decline in the number of engineering graduates \n",
    "    and a lack of well-educated engineers.\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'generated_text': 'This laboratory was adapted from the original course by HuggingFace.'}]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Translation: Translate a sentence from French to English\n",
    "# The same result can be achieved using the `translation` pipeline.\n",
    "from transformers import pipeline\n",
    "\n",
    "translator = pipeline(\"text2text-generation\", model=\"Helsinki-NLP/opus-mt-fr-en\")\n",
    "translator(\"Ce laboratoire a été adapté à partir du cours originel par HuggingFace.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This concludes our pipeline overview. Refer to the [documentation](https://huggingface.co/docs/transformers/v4.15.0/en/main_classes/pipelines) for additional information on pipeline types and parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Behind the Pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we saw in the previous chapter, a `pipeline` has a preprocessing step, a model inference step and a post-processing step:\n",
    "\n",
    "<div>\n",
    "<img src=\"https://huggingface.co/course/static/chapter2/full_nlp_pipeline.png\", alt=\"Visual representation of a full NLP pipeline\" width=\"80%\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like other neural networks, Transformer models can’t process raw text directly, so the first step of our pipeline is to convert the text inputs into numbers that the model can make sense of. To do this we use a tokenizer, which will be responsible for:\n",
    "\n",
    "- Splitting the input into words, subwords, or symbols (like punctuation) that are called tokens\n",
    "- Mapping each token to an integer\n",
    "- Adding additional inputs that may be useful to the model\n",
    "\n",
    "All this preprocessing needs to be done in exactly the same way as when the model was pretrained, so we first need to download that information from the Model Hub. To do this, we use the `AutoTokenizer` class and its `from_pretrained` method. Using the checkpoint name of our model, it will automatically fetch the data associated with the model’s tokenizer and cache it (so it’s only downloaded the first time you run the code below).\n",
    "\n",
    "**Important:** Every type of transformer model has its own tokenizer and model classes (e.g. `BertTokenizer`, `GPT2Tokenizer`, `XLMTokenizer`, ...). The `AutoTokenizer` and `AutoModel` classess will rely on configurations saved alongside model checkpoints to load any model in the right class, so you don’t need to worry about this.\n",
    "\n",
    "Let's try to load and use the tokenizer of the sentiment analysis model we tried at the beginning of this lab:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vocabulary size: 30522 tokens\n",
      "10 random tokens: {'chesapeake': 20867, 'incorporating': 13543, 'acted': 6051, 'prints': 11204, '##arium': 17285, 'ventura': 21151, 'groves': 21695, 'patent': 7353, '[unused699]': 704, 'nec': 26785}\n",
      "==================== \n",
      " Output: {'input_ids': tensor([[ 101, 1045, 1005, 2310, 2042, 3403, 2005, 2023, 2607, 2005, 2026, 2878,\n",
      "         2166, 1012,  102],\n",
      "        [ 101, 1045, 5223, 2023, 2607, 2061, 2172,  999,  102,    0,    0,    0,\n",
      "            0,    0,    0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
      "        [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]])}\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "identifier = \"distilbert-base-uncased-finetuned-sst-2-english\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(identifier)\n",
    "\n",
    "print(\"Vocabulary size:\", len(tokenizer), \"tokens\")\n",
    "print(\"10 random tokens:\", {k:v for i, (k,v) in enumerate(tokenizer.vocab.items()) if i < 10})\n",
    "\n",
    "raw_inputs = [\"I've been waiting for this course for my whole life.\", \"I hate this course so much!\"]\n",
    "inputs = tokenizer(raw_inputs, padding=True, return_tensors=\"pt\")\n",
    "print(\"=\"*20,\"\\n\",\"Output:\",inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `padding` argument tells the tokenizer to make sure all input sequences are encoded with the same length. This is important since the model needs batches of IDs having the same size to work properly. The `return_tensors` argument tells the tokenizer to return the output as a PyTorch tensor instead of a list of tokens. This is important since the model expects the input to be a tensor.\n",
    "\n",
    "The output itself is a dictionary containing two keys, `input_ids` and `attention_mask`. `input_ids` contains two rows of integers (one for each sentence) that are the unique identifiers of the tokens in each sentence. The `attention_mask` is used by the model to ignore padding tokens when computing the loss.\n",
    "\n",
    "We can recover the original tokens from input ids as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['[CLS]', 'i', \"'\", 've', 'been', 'waiting', 'for', 'this', 'course', 'for', 'my', 'whole', 'life', '.', '[SEP]']\n",
      "['[CLS]', 'i', 'hate', 'this', 'course', 'so', 'much', '!', '[SEP]', '[PAD]', '[PAD]', '[PAD]', '[PAD]', '[PAD]', '[PAD]']\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.convert_ids_to_tokens(inputs[\"input_ids\"][0]))\n",
    "print(tokenizer.convert_ids_to_tokens(inputs[\"input_ids\"][1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The special `[CLS]` and `[SEP]` tokens are used by BERT-like models to delimit the sentences, and they are automatically added by the tokenizer.\n",
    "\n",
    "We can now proceed to download the model and perform inference over the input ids:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilbert-base-uncased-finetuned-sst-2-english were not used when initializing DistilBertModel: ['classifier.bias', 'pre_classifier.bias', 'classifier.weight', 'pre_classifier.weight']\n",
      "- This IS expected if you are initializing DistilBertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DistilBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModel\n",
    "\n",
    "identifier = \"distilbert-base-uncased-finetuned-sst-2-english\"\n",
    "model = AutoModel.from_pretrained(identifier)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This architecture contains only the base Transformer module: given some inputs, it outputs what we’ll call **hidden states**, also known as features. For each model input, we’ll retrieve a high-dimensional vector representing the contextual representation of that token in the model's learned embedding space.\n",
    "\n",
    "While these hidden states can be useful on their own, they’re usually inputs to another part of the model, known as the **head**, which is responsible to perform the actual prediction associated to the target task.\n",
    "\n",
    "The input vector to the model is usually three-dimensional, containing respectively\n",
    "\n",
    "- Batch size: The number of sequences processed at a time (2 in our example).\n",
    "- Sequence length: The length of the numerical representation of the sequence (15 in our example).\n",
    "- Hidden size: The vector dimension of each model input. The vector is said to be “high dimensional” because of this last value. The hidden size can be very large (768 is common for smaller models, and in larger models this can reach 3072 or more).\n",
    "\n",
    "We can now feed the `inputs` produced by the tokenizer to the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 15, 768])\n"
     ]
    }
   ],
   "source": [
    "outputs = model(**inputs)\n",
    "print(outputs.last_hidden_state.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are the same values that could be extracted by the `feature-extraction` pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using Heads"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model heads take the high-dimensional vector of hidden states as input and project them onto a different dimension. They are usually composed of one or a few linear layers:\n",
    "\n",
    "<div>\n",
    "<img src=\"https://huggingface.co/course/static/chapter2/transformer_and_head.png\", alt=\"A Transformer model\" width=\"80%\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output of the Transformer model is sent directly to the model head to be processed.\n",
    "\n",
    "In this diagram, the model is represented by its embeddings layer and the subsequent layers. The embeddings layer converts each input ID in the tokenized input into a vector that represents the associated token. The subsequent layers manipulate those vectors using the attention mechanism to produce the final representation of the sentences.\n",
    "\n",
    "There are many different architectures available in 🤗 Transformers, with each one designed around tackling a specific task. Here is a non-exhaustive list:\n",
    "\n",
    "- *Model (default headless model)\n",
    "- *ForCausalLM\n",
    "- *ForMaskedLM\n",
    "- *ForMultipleChoice\n",
    "- *ForQuestionAnswering\n",
    "- *ForSequenceClassification\n",
    "- *ForTokenClassification\n",
    "- *ForSeq2SeqLM\n",
    "\n",
    "For our example, we will need a model with a sequence classification head (to be able to classify the sentences as positive or negative). So, we won’t actually use the `AutoModel` class, but `AutoModelForSequenceClassification`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSequenceClassification\n",
    "\n",
    "checkpoint = \"distilbert-base-uncased-finetuned-sst-2-english\"\n",
    "model = AutoModelForSequenceClassification.from_pretrained(checkpoint)\n",
    "outputs = model(**inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now if we look at the shape of our inputs, the dimensionality will be much lower: the model head takes as input the high-dimensional vectors we saw before, and outputs vectors containing two values (one per label). Since we have just two sentences and two labels, the result we get from our model is of shape 2 x 2:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 2])\n"
     ]
    }
   ],
   "source": [
    "print(outputs.logits.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Postprocessing outputs\n",
    "\n",
    "The values we get as output from our model don’t necessarily make sense by themselves. Let’s take a look:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-2.0023,  1.9307],\n",
      "        [ 4.4057, -3.5342]], grad_fn=<AddmmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "print(outputs.logits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our model predicted [-2.0023, 1.9307] for the first sentence and [4.4057, -3.5342] for the second one. Those are not probabilities but logits, the raw, unnormalized scores outputted by the last layer of the model. To be converted to probabilities, they need to go through a SoftMax layer (all 🤗 Transformers models output the logits, as the loss function for training will generally fuse the last activation function, such as SoftMax, with the actual loss function, such as cross entropy):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1.9210e-02, 9.8079e-01],\n",
      "        [9.9964e-01, 3.5611e-04]], grad_fn=<SoftmaxBackward0>)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)\n",
    "print(predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can see that the model predicted [0.0192, 0.9807] for the first sentence and [0.9996, 0.0004] for the second one. These are recognizable probability scores.\n",
    "\n",
    "To get the labels corresponding to each position, we can inspect the `id2label` attribute of the model configuration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 'NEGATIVE', 1: 'POSITIVE'}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.config.id2label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can conclude that the model predicted the following:\n",
    "\n",
    "- First sentence: NEGATIVE: 0.0192, POSITIVE: 0.9807\n",
    "- Second sentence: NEGATIVE: 0.9996, POSITIVE: 0.0004\n",
    "\n",
    "We have successfully reproduced the three steps of the pipeline: preprocessing with tokenizers, passing the inputs through the model, and postprocessing. Here is a summary of all the steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I've been waiting for this course for my whole life. {'NEGATIVE': 0.019209884107112885, 'POSITIVE': 0.9807901382446289}\n",
      "I hate this course so much! {'NEGATIVE': 0.9996439218521118, 'POSITIVE': 0.0003561103658284992}\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "\n",
    "identifier = \"distilbert-base-uncased-finetuned-sst-2-english\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(identifier)\n",
    "model = AutoModelForSequenceClassification.from_pretrained(identifier)\n",
    "sequences = [\"I've been waiting for this course for my whole life.\", \"I hate this course so much!\"]\n",
    "\n",
    "\n",
    "tokens = tokenizer(sequences, padding=True, return_tensors=\"pt\")\n",
    "output = model(**tokens)\n",
    "probs = torch.nn.functional.softmax(output.logits, dim=-1).tolist()\n",
    "for i, p in enumerate(probs):\n",
    "    print(sequences[i], {model.config.id2label[j]:v for j, v in enumerate(p)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the case of generative models (both causal like GPT-2 and sequence-to-sequence like BART or T5), we make use of the `model.generate` method to generate a sequence of ids representing the predicted output. We then use the `tokenizer.decode` method to convert the ids into the corresponding text. We will use a small transformer model trained to translate from English to Russian as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['My name is Gabriele and I live in Groningen.', 'I hope you like the course!']\n",
      "My name is Gabriele and I live in Groningen.\n",
      "I hope you like the course!\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n",
    "\n",
    "identifier = \"Helsinki-NLP/opus-mt-ru-en\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(identifier)\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(identifier)\n",
    "sequences = [\"Меня зовут Габриелэ и я живу в Гроннингене\", \"Надеюсь, вам нравится курс!\"]\n",
    "\n",
    "\n",
    "tokens = tokenizer(sequences, padding=True, return_tensors=\"pt\")\n",
    "output = model.generate(**tokens)\n",
    "print(tokenizer.batch_decode(output, skip_special_tokens=True))\n",
    "\n",
    "# Alternatively, use decode sentence-by-sentence\n",
    "for s in output.tolist():\n",
    "    print(tokenizer.decode(s, skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `skip_special_token` directive removes all the special tokens like `</s>` and `<pad>` from the output. The `generate` method is highly versatile, refer to the [documentation](https://huggingface.co/docs/transformers/v4.15.0/en/main_classes/model#transformers.generation_utils.GenerationMixin.generate) for more details. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading and Saving"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As mentioned before, the `AutoModel` class and all of its relatives are actually simple wrappers over the wide variety of models available in the library. It’s a clever wrapper as it can automatically guess the appropriate model architecture for your checkpoint, and then instantiates a model with this architecture. However, if you know the type of model you want to use, you can use the class that defines its architecture directly. Let’s take a look at how this works with a BERT model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BertConfig {\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"classifier_dropout\": null,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"layer_norm_eps\": 1e-12,\n",
      "  \"max_position_embeddings\": 512,\n",
      "  \"model_type\": \"bert\",\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"pad_token_id\": 0,\n",
      "  \"position_embedding_type\": \"absolute\",\n",
      "  \"transformers_version\": \"4.12.5\",\n",
      "  \"type_vocab_size\": 2,\n",
      "  \"use_cache\": true,\n",
      "  \"vocab_size\": 30522\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from transformers import BertConfig, BertModel\n",
    "\n",
    "# Building the config\n",
    "config = BertConfig()\n",
    "\n",
    "# Building the model from the config\n",
    "model = BertModel(config)\n",
    "\n",
    "print(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The configuration contains many attributes that are used to build the model. While you haven’t seen what all of these attributes do yet, you should recognize some of them: the `hidden_size` attribute defines the size of the `hidden_states` vector, and `num_hidden_layers` defines the number of layers the Transformer model has.\n",
    "\n",
    "Creating a model from the default configuration initializes it with random values. The model can be used in this state, but it will output gibberish; it needs to be trained first. However, this procedure requires a long time and a lot of data. To avoid unnecessary and duplicated effort, it’s imperative to be able to share and reuse models that have already been trained.\n",
    "\n",
    "Loading a Transformer model that is already trained is simple — we can do this using the `from_pretrained` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertModel\n",
    "\n",
    "model = BertModel.from_pretrained(\"bert-base-cased\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This model is now initialized with all the weights of the checkpoint. It can be used directly for inference on the tasks it was trained on, and it can also be fine-tuned on a new task. By training with pretrained weights rather than from scratch, we can quickly achieve good results.\n",
    "\n",
    "The weights have been downloaded and cached (so future calls to the `from_pretrained` method won’t re-download them) in the cache folder, which defaults to ~/.cache/huggingface/transformers. You can customize your cache folder by setting the HF_HOME environment variable.\n",
    "\n",
    "The identifier used to load the model can be the identifier of any model on the Model Hub, as long as it is compatible with the BERT architecture.\n",
    "\n",
    "Saving a model is as easy as loading one — we use the `save_pretrained` method, which is analogous to the `from_pretrained` method. This saves two files to your disk: the configuration file and the weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(\"directory_on_my_computer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same can be accomplished for a tokenizer. This saves the essential files to restore the tokenizer object using `from_pretrained(directory_on_my_computer)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('directory_on_my_computer/tokenizer_config.json',\n",
       " 'directory_on_my_computer/special_tokens_map.json',\n",
       " 'directory_on_my_computer/vocab.txt',\n",
       " 'directory_on_my_computer/added_tokens.json')"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import BertTokenizer\n",
    "\n",
    "tokenizer = BertTokenizer.from_pretrained(\"bert-base-cased\")\n",
    "tokenizer.save_pretrained(\"directory_on_my_computer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This concludes our quick overview of the `transformers` library. You can refer to the extended version of this introduction at the original [HuggingFace Course](https://huggingface.co/course) for more examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the 🤗 Datasets library"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "🤗 Datasets provides loading scripts to handle the loading of local and remote datasets. It supports several common data formats, such as:\n",
    "\n",
    "Data format \t   |Loading script |Example\n",
    "-------------------|---------------|----------------------------------------------|\n",
    "CSV & TSV \t       |csv \t       |`load_dataset(\"csv\", data_files=\"my_file.csv\")`|\n",
    "Text files \t       |text \t       |`load_dataset(\"text\", data_files=\"my_file.txt\")`|\n",
    "JSON & JSON Lines  |json \t       |`load_dataset(\"json\", data_files=\"my_file.jsonl\")`|\n",
    "Pickled DataFrames |pandas \t       |`load_dataset(\"pandas\", data_files=\"my_dataframe.pkl\")`|\n",
    "\n",
    "As shown in the table, for each data format we just need to specify the type of loading script in the `load_dataset` function, along with a data_files argument that specifies the path to one or more files. Let’s see how we can load some remote files from the [SQUAD-it](https://github.com/crux82/squad-it) dataset for question answering in Italian:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration default-57dcee3ea6992346\n",
      "Reusing dataset json (/home/gsarti/.cache/huggingface/datasets/json/default-57dcee3ea6992346/0.0.0/c2d554c3377ea79c7664b93dc65d0803b45e3279000f993c7bfd18937fd7f426)\n",
      "100%|██████████| 2/2 [00:00<00:00, 103.16it/s]\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "url = \"https://github.com/crux82/squad-it/raw/master/\"\n",
    "data_files = {\n",
    "    \"train\": url + \"SQuAD_it-train.json.gz\",\n",
    "    \"test\": url + \"SQuAD_it-test.json.gz\",\n",
    "}\n",
    "squad_it_dataset = load_dataset(\"json\", data_files=data_files, field=\"data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the dataset contains two splits, `train` and `test`, which contain respectively 442 and 48 pairs of (title, paragraph) fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['title', 'paragraphs'],\n",
       "        num_rows: 442\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['title', 'paragraphs'],\n",
       "        num_rows: 48\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squad_it_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now inspect the first pair, showing that it contains a `paragraph` and a nested set of `qas`. To do so, we access the `train` dataset in the DatasetDict as we would access a normal dictionary, and select the first element from it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "squad_it_dataset[\"train\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now load the original [SQUAD](https://rajpurkar.github.io/SQuAD-explorer/) dataset directly from the Datasets hub, using its identified `squad`. We can see that the structure has been flattened and we have now a triplet (context, question, answer) per row, for a total of 87599 training examples and 10570 validation examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reusing dataset squad (/home/gsarti/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453)\n",
      "100%|██████████| 2/2 [00:00<00:00, 162.40it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'title', 'context', 'question', 'answers'],\n",
       "        num_rows: 87599\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['id', 'title', 'context', 'question', 'answers'],\n",
       "        num_rows: 10570\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "squad = load_dataset(\"squad\")\n",
    "squad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get a feel of the dataset, let's grab a sample of two examples selected at random from the `train` split. The result of the slicing operation is a dictionary with lists containing the values for fields for each one of the examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached shuffled indices for dataset at /home/gsarti/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453/cache-10de9997c4b83f65.arrow\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'id': ['573173d8497a881900248f0c', '57277e815951b619008f8b52'],\n",
       " 'title': ['Egypt', 'Ann_Arbor,_Michigan'],\n",
       " 'context': ['The Pew Forum on Religion & Public Life ranks Egypt as the fifth worst country in the world for religious freedom. The United States Commission on International Religious Freedom, a bipartisan independent agency of the US government, has placed Egypt on its watch list of countries that require close monitoring due to the nature and extent of violations of religious freedom engaged in or tolerated by the government. According to a 2010 Pew Global Attitudes survey, 84% of Egyptians polled supported the death penalty for those who leave Islam; 77% supported whippings and cutting off of hands for theft and robbery; and 82% support stoning a person who commits adultery.',\n",
       "  'The Ann Arbor Hands-On Museum is located in a renovated and expanded historic downtown fire station. Multiple art galleries exist in the city, notably in the downtown area and around the University of Michigan campus. Aside from a large restaurant scene in the Main Street, South State Street, and South University Avenue areas, Ann Arbor ranks first among U.S. cities in the number of booksellers and books sold per capita. The Ann Arbor District Library maintains four branch outlets in addition to its main downtown building. The city is also home to the Gerald R. Ford Presidential Library.'],\n",
       " 'question': ['What percentage of Egyptians polled support death penalty for those leaving Islam?',\n",
       "  'Ann Arbor ranks 1st among what goods sold?'],\n",
       " 'answers': [{'text': ['84%'], 'answer_start': [468]},\n",
       "  {'text': ['books'], 'answer_start': [402]}]}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squad_train = squad[\"train\"].shuffle(seed=42).select(range(1000))\n",
    "squad_train[:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may have noticed that the field `answers` contains a list of dictionaries, each one containing a list of `answer_start` and a list of `text` fields. This is because the dataset contains multiple answers per question, and because extractive QA systems use the `answer_start` field as prediction target.\n",
    "\n",
    "Let's say we plan to train a sequence-to-sequence model to generate answers from paragraphs and questions, without extracting them directly from the context. In this case, the index of the answer in the `answers` field is not relevant, and we can remove it. We want to adapt the dataset by extracting only the value of the `text` field from the first answer for each example, and by concatenating the context with the question to have a single source field. Let's define a custom function and use the `Dataset.map` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/gsarti/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453/cache-b73576d9a85a1624.arrow\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'id': ['573173d8497a881900248f0c', '57277e815951b619008f8b52'],\n",
       " 'text_answer': ['84%', 'books'],\n",
       " 'text_question': ['The Pew Forum on Religion & Public Life ranks Egypt as the fifth worst country in the world for religious freedom. The United States Commission on International Religious Freedom, a bipartisan independent agency of the US government, has placed Egypt on its watch list of countries that require close monitoring due to the nature and extent of violations of religious freedom engaged in or tolerated by the government. According to a 2010 Pew Global Attitudes survey, 84% of Egyptians polled supported the death penalty for those who leave Islam; 77% supported whippings and cutting off of hands for theft and robbery; and 82% support stoning a person who commits adultery. Question: What percentage of Egyptians polled support death penalty for those leaving Islam?',\n",
       "  'The Ann Arbor Hands-On Museum is located in a renovated and expanded historic downtown fire station. Multiple art galleries exist in the city, notably in the downtown area and around the University of Michigan campus. Aside from a large restaurant scene in the Main Street, South State Street, and South University Avenue areas, Ann Arbor ranks first among U.S. cities in the number of booksellers and books sold per capita. The Ann Arbor District Library maintains four branch outlets in addition to its main downtown building. The city is also home to the Gerald R. Ford Presidential Library. Question: Ann Arbor ranks 1st among what goods sold?']}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def extract_text(example):\n",
    "    return {\n",
    "        \"text_answer\": example[\"answers\"][\"text\"][0],\n",
    "        \"text_question\": example[\"context\"] + \" Question: \" + example[\"question\"],\n",
    "    }\n",
    "\n",
    "squad_train = squad_train.map(\n",
    "    extract_text, remove_columns=[\"title\", \"context\", \"question\", \"answers\"]\n",
    ")\n",
    "\n",
    "squad_train[:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see we are using the `remove_columns` parameter in `map` to also drop unused columns. Let's give the last touches by filtering out examples with context exceeding 512 words, and renaming the two new fields to `source` and `target` for uniformity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/gsarti/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453/cache-f61fc212407c59e3.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length original sample: 1000 Length filtered sample: 1000\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'id': ['573173d8497a881900248f0c', '57277e815951b619008f8b52'],\n",
       " 'target': ['84%', 'books'],\n",
       " 'source': ['The Pew Forum on Religion & Public Life ranks Egypt as the fifth worst country in the world for religious freedom. The United States Commission on International Religious Freedom, a bipartisan independent agency of the US government, has placed Egypt on its watch list of countries that require close monitoring due to the nature and extent of violations of religious freedom engaged in or tolerated by the government. According to a 2010 Pew Global Attitudes survey, 84% of Egyptians polled supported the death penalty for those who leave Islam; 77% supported whippings and cutting off of hands for theft and robbery; and 82% support stoning a person who commits adultery. Question: What percentage of Egyptians polled support death penalty for those leaving Islam?',\n",
       "  'The Ann Arbor Hands-On Museum is located in a renovated and expanded historic downtown fire station. Multiple art galleries exist in the city, notably in the downtown area and around the University of Michigan campus. Aside from a large restaurant scene in the Main Street, South State Street, and South University Avenue areas, Ann Arbor ranks first among U.S. cities in the number of booksellers and books sold per capita. The Ann Arbor District Library maintains four branch outlets in addition to its main downtown building. The city is also home to the Gerald R. Ford Presidential Library. Question: Ann Arbor ranks 1st among what goods sold?']}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# We use a lambda expression here for simple whitespace tokenization, but a function\n",
    "# using an AutoTokenizer could also be used.\n",
    "filtered_train = squad_train.filter(lambda x: len(x[\"text_question\"].split(\" \")) < 512)\n",
    "print(\"Length original sample:\", len(squad_train), \"Length filtered sample:\", len(filtered_train))\n",
    "\n",
    "filtered_train = filtered_train.rename_columns({\n",
    "    \"text_answer\": \"target\",\n",
    "    \"text_question\": \"source\"\n",
    "})\n",
    "\n",
    "filtered_train[:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To enable the conversion between various third-party libraries, 🤗 Datasets provides a `Dataset.set_format` function. This function only changes the output format of the dataset, so you can easily switch to another format without affecting the underlying data format, which is Apache Arrow. The formatting is done in place. To demonstrate, let’s convert our dataset to the popular Pandas library used in data science:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>target</th>\n",
       "      <th>source</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>573173d8497a881900248f0c</td>\n",
       "      <td>84%</td>\n",
       "      <td>The Pew Forum on Religion &amp; Public Life ranks ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>57277e815951b619008f8b52</td>\n",
       "      <td>books</td>\n",
       "      <td>The Ann Arbor Hands-On Museum is located in a ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>5727e2483acd2414000deef0</td>\n",
       "      <td>the executive</td>\n",
       "      <td>One important aspect of the rule-of-law initia...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                         id         target  \\\n",
       "0  573173d8497a881900248f0c            84%   \n",
       "1  57277e815951b619008f8b52          books   \n",
       "2  5727e2483acd2414000deef0  the executive   \n",
       "\n",
       "                                              source  \n",
       "0  The Pew Forum on Religion & Public Life ranks ...  \n",
       "1  The Ann Arbor Hands-On Museum is located in a ...  \n",
       "2  One important aspect of the rule-of-law initia...  "
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Can be reset using reset_format()\n",
    "filtered_train.set_format(\"pandas\")\n",
    "filtered_train[:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create a proper DataFrame from the dataset, and perform some additional transformations on it to see what's the average length of answers in words:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>num_words</th>\n",
       "      <th>frequency</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>341</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>246</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>168</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>77</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>42</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   num_words  frequency\n",
       "0          1        341\n",
       "1          2        246\n",
       "2          3        168\n",
       "3          4         77\n",
       "4          5         42"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df = filtered_train[:]\n",
    "train_df[\"len_target\"] = train_df[\"target\"].apply(lambda x: len(x.split(\" \")))\n",
    "\n",
    "frequencies = (\n",
    "    train_df[\"len_target\"]\n",
    "    .value_counts()\n",
    "    .to_frame()\n",
    "    .reset_index()\n",
    "    .rename(columns={\"index\": \"num_words\", \"len_target\": \"frequency\"})\n",
    ")\n",
    "frequencies.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now retransform this new DataFrame into a Dataset object using the `Dataset.from_pandas` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['num_words', 'frequency'],\n",
       "    num_rows: 26\n",
       "})"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import Dataset\n",
    "\n",
    "freq_dataset = Dataset.from_pandas(frequencies)\n",
    "freq_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To conclude this overview of the 🤗 Datasets library, let's take a look at how Datasets can be saved. The three main functions to save a Dataset are `save_to_disk`, `to_csv` and `to_json`. \n",
    "\n",
    "The first one creates a folder where we can see that each split is associated with its own dataset.arrow table, and some metadata in dataset_info.json and state.json. You can think of the Arrow format as a fancy table of columns and rows that is optimized for building high-performance applications that process and transport large datasets.\n",
    "\n",
    "Once the dataset is saved, it can be loaded back using `load_from_disk`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['num_words', 'frequency'],\n",
       "    num_rows: 26\n",
       "})"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_from_disk\n",
    "\n",
    "# Save all splits to disk\n",
    "freq_dataset.save_to_disk(\"frequencies\")\n",
    "del freq_dataset\n",
    "\n",
    "# Reload them into a single DatasetDict object\n",
    "freq_dataset = load_from_disk(\"frequencies\")\n",
    "freq_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the CSV and JSON formats, we have to store each split as a separate file. One way to do this is by iterating over the keys and values in the DatasetDict object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Creating json from Arrow format: 100%|██████████| 9/9 [00:00<00:00, 14.15ba/s]\n",
      "Creating json from Arrow format: 100%|██████████| 2/2 [00:00<00:00, 22.97ba/s]\n"
     ]
    }
   ],
   "source": [
    "for split, dataset in squad.items():\n",
    "    dataset.to_json(f\"squad-{split}.jsonl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then reload the dataset using the `load_dataset` function, as seen previously:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration default-260a0a433a957006\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset json/default to /home/gsarti/.cache/huggingface/datasets/json/default-260a0a433a957006/0.0.0/c2d554c3377ea79c7664b93dc65d0803b45e3279000f993c7bfd18937fd7f426...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00, 1203.70it/s]\n",
      "100%|██████████| 2/2 [00:00<00:00, 1218.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset json downloaded and prepared to /home/gsarti/.cache/huggingface/datasets/json/default-260a0a433a957006/0.0.0/c2d554c3377ea79c7664b93dc65d0803b45e3279000f993c7bfd18937fd7f426. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00, 301.62it/s]\n"
     ]
    }
   ],
   "source": [
    "data_files = {\n",
    "    \"train\": \"squad-train.jsonl\",\n",
    "    \"validation\": \"squad-validation.jsonl\",\n",
    "}\n",
    "squad = load_dataset(\"json\", data_files=data_files)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This concludes our quick overview of the 🤗 Datasets library. For more advanced concepts related to Datasets (memory mapping, interleaving, streaming, semantic search) refer to the Chapters 5.4 and 5.5 of the [HuggingFace Course](https://huggingface.co/course/chapter5/4?fw=pt)."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "f9f85d98b63f393548f3009c8d52d8286e609a1467b1184fe464fb700873fbd3"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit ('venv': venv)",
   "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.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}