{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dbd51d48",
   "metadata": {},
   "source": [
    "# MAGMA Inference Demo\n",
    "\n",
    "Copyright (c) 2023 Graphcore Ltd."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de6ceb61",
   "metadata": {},
   "source": [
    "This notebook provides a basic interactive MAGMA inference application, based on the freely available checkpoint published by Aleph Alpha. Note that such a checkpoint is only a demo meant to help users understand how the model works."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90bd2cb1",
   "metadata": {},
   "source": [
    "MAGMA (Multimodal Augmentation of Generative Models through Adapter-based Finetuning) is a multimodal Vision-Language model developed by Aleph Alpha and researchers from [Heidelberg University](https://www.cl.uni-heidelberg.de \"Computational Linguistics at Heidelberg University\"). For all the details, please refer to the [research paper](https://arxiv.org/abs/2112.05253) and the [official GitHub repository](https://github.com/Aleph-Alpha/magma). The MAGMA model autoregressively generates text from arbitrary combinations of visual and textual input."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e68addf8",
   "metadata": {},
   "source": [
    "MAGMA aims at obtaining a joint image-text representation, which can then be used for a variety of tasks such as image captioning, visual question answering and visual entailment. You can feed the model with a textual prompt and an image and use the textual prompt to make queries about the image.\n",
    "\n",
    "The main components of MAGMA are:\n",
    "- A pretrained autoregressive **language model**. This component is **frozen**, it is not trained. The chosen language model is [GPT-J 6B](https://github.com/kingoflolz/mesh-transformer-jax).\n",
    "- A **visual encoder** (namely, variants of ViT or ResNet) that takes an image input and produces image features. This is a **trainable** component. In the paper, several variants are discussed and compared.\n",
    "- An **image prefix** that projects image features\n",
    "into a sequence of embedding vectors, suitable to be inputs of the language model transformer. This is a **trainable** component.\n",
    "- **Adapter** layers, which are **trainable** layers added in different layouts to the language model decoder block. In the paper, several configurations are discussed and compared.\n",
    "\n",
    "The image below, taken from the Aleph Alpha [paper](https://arxiv.org/abs/2112.05253), summarises the model structure."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73d2c47f",
   "metadata": {},
   "source": [
    "![MAGMA model architecture](images/MagmaStructure.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b94fdd93",
   "metadata": {},
   "source": [
    "- The inputs of the model are an image and an accompanying text.\n",
    "- Text is encoded by the language model embedding **E**, obtaining text embeddings.\n",
    "- In the ImagePrefix:\n",
    "    - Image features are produced by the visual encoder **V<sup>e</sup>**\n",
    "    - Image features are projected to the embedding dimension by a linear layer **V<sup>p</sup>**, obtaining image embeddings.\n",
    "- Image embeddings and text embeddings are concatenated. Concatenation order is `(image, text)`, with the image tokens always at the beginning of the sequence. In this way, they are always attended by the text tokens (not affected by causal mask).\n",
    "- The resulting sequence is fed to the language model transformer endowed with adapters.\n",
    "\n",
    "This application implements only the model choices of the [publicly available checkpoint](https://bit.ly/aleph_alpha_magma_download). As they state in the model repo, this checkpoint is just meant to be a demo and it's not the one they use in actual applications:\n",
    "- The visual encoder is a CLIP modified ResNet50-x16 ([implementation](https://github.com/openai/CLIP/blob/main/clip/model.py)), without the final attention pool layer.\n",
    "- Adapters are only added to feed-forward blocks. In the MAGMA paper, it says that adding adapters also to the attention layers gives better results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5847a47a",
   "metadata": {},
   "source": [
    "## Inference\n",
    "\n",
    "**Note:** Magma uses GPT-J, but in the code you can find references to GPTNeo, and this can be confusing. The model is GPT-J 6B but not taken from the official Hugging Face transformer library. Instead, they use [finetuneanon/transformer](https://github.com/finetuneanon/transformers). In this library, [GPTNeo](https://github.com/finetuneanon/transformers/tree/gpt-neo-localattention3-rp-b/src/transformers/models/gpt_neo) with the options `jax=True`, `rotary=True` and `global` attention corresponds to GPT-J.\n",
    "\n",
    "### Input preprocessing\n",
    "#### Textual input\n",
    "The input text is tokenized using GPT2 tokenizer.\n",
    "It is then padded up to `sequence_len - 144`, where `sequence_len` is specified in the configs and 144 is the number of image tokens generated by the image encoder.\n",
    "#### Image input\n",
    "The input image is resized, cropped and normalised using Magma [clip_preprocess](https://github.com/Aleph-Alpha/magma/blob/master/magma/transforms.py) function.\n",
    "\n",
    "\n",
    "### Next token generation\n",
    "Given an image and a text prompt, the model outputs logits for the sequence.\n",
    "Logits corresponding to the last token in the sequence are used to predict the next token.\n",
    "Several heuristics are supported:\n",
    "- argmax: simply pick the token with the highest probability. This is a deterministic sampling method.\n",
    "- top-p: probability is redistributed among the first x tokens such that the cumulative probability is greater than a threshold p. Then, next token is sampled from such distribution (categorical sampling, non deterministic).\n",
    "- top-k: probability is redistributed among the K most-likely tokens. Then, next token is sampled from such distribution (categorical sampling, non deterministic).\n",
    "- temperature: logits are scaled by a factor 1/T (T between 0 and 1 ) before applying the softmax. This makes the distribution more peaked for low temperature, and broader for high temperatures. A zero temperature corresponds to a deterministic choice (argmax), while sampling output becomes more random as we increase the temperature.\n",
    "\n",
    "If you are not familiar with these concepts, this [HuggingFace article](https://huggingface.co/blog/how-to-generate) can help you visualise them.\n",
    "\n",
    "The new token is added to the input sequence, and the process goes on until `max_out_tokens` are generated or the model outputs the `EOS` token.\n",
    "\n",
    "### Execution scheme\n",
    "The model is run using phased execution.\n",
    "This means that the model is partitioned into a set of smaller graphs that are executed in series on the IPU, using remote memory to store variables and input/output tensors between calls (activations).\n",
    "We recommend going through the tutorial [Phased Execution in MNIST example](https://github.com/graphcore/examples/tree/master/tutorials/tutorials/popxl/6_phased_execution) to better understand this execution scheme.\n",
    "- The first phase corresponds to the image prefix. The output of this phase produces image embeddings that are concatenated to the (tokenized) textual input\n",
    "- The following phase is the GPT-J embedding phase, which produces text embeddings.\n",
    "- Image embeddings and text embeddings are concatenated and fed to GPT-J blocks. Each block constitutes a different phase.\n",
    "- The final phase is the LM head, producing logits to be used for the next token prediction.\n",
    "\n",
    "The GPT-J model makes use of tensor parallelism, spanning across 4 IPUs. More details about the implementation are available in the [GPT-J Notebook]().\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49ea7cc6",
   "metadata": {},
   "source": [
    "## Environment setup\n",
    "\n",
    "The best way to run this demo is on Paperspace Gradient’s cloud IPUs because everything is already set up for you. To improve your experience, we preload datasets and pre-install packages. This can take a few minutes. If you experience errors immediately after starting a session, please try restarting the kernel before contacting support. If a problem persists or you want to give us feedback on the content of this notebook, please reach out to through our community of developers using our [Slack channel](https://www.graphcore.ai/join-community). \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e83cc038",
   "metadata": {},
   "source": [
    "To run the demo using other IPU hardware, you need to have the Poplar SDK enabled. Refer to the [Getting Started guide](https://docs.graphcore.ai/en/latest/getting-started.html#getting-started) for your system for details on how to do this. Also refer to the [Jupyter Quick Start guide](https://docs.graphcore.ai/projects/jupyter-notebook-quick-start/en/latest/index.html) for how to set up Jupyter to be able to run this notebook on a remote IPU machine.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "146240fb",
   "metadata": {},
   "source": [
    "In order to improve usability and support for future users, Graphcore would like to collect information about the applications and code being run in this notebook. The following information will be anonymised before being sent to Graphcore:\n",
    "\n",
    "User progression through the notebook\n",
    "Notebook details: number of cells, code being run and the output of the cells\n",
    "Environment details\n",
    "You can disable logging at any time by running %unload_ext graphcore_cloud_tools.notebook_logging.gc_logger from any cell."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66d52b7d",
   "metadata": {},
   "source": [
    "Uncomment the cell below to install libopenmpi-dev if not available in your docker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdab6495",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!export DEBIAN_FRONTEND=noninteractive\n",
    "#!apt-get update\n",
    "#!apt install -y libopenmpi-dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09d9bc31",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt --ignore-requires-python\n",
    "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58a6569b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from run_inference import run_inference, init_inference_session\n",
    "import sys, os, os.path\n",
    "\n",
    "executable_cache_dir = os.path.join(\n",
    "    os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"./exe_cache\"), \"magma\"\n",
    ")\n",
    "os.environ[\"POPXL_CACHE_DIR\"] = executable_cache_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04412adc",
   "metadata": {},
   "source": [
    "## Compile and load the model\n",
    "Compilation takes around 3 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f8aeb56-0b49-418d-942c-ec15e0542ed5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from magma.magma import Magma\n",
    "from configs import CONFIG_DIR, MagmaConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45687ddf",
   "metadata": {},
   "source": [
    "The upstream `Magma.from_checkpoint` which is called in `init_inference_session()`` uses `gdown`, which appears to be broken at the moment. Either download the checkpoint manually, or allow `wget` to download it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11176741-8d43-4b8d-a8c9-6d66d72b8542",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_path = os.path.join(\n",
    "    os.getenv(\"DATASETS_DIR\", \"./\"), \"mp_rank_00_model_states.pt\"\n",
    ")\n",
    "\n",
    "if os.path.exists(checkpoint_path) == False:\n",
    "    !wget 'https://bit.ly/aleph-alpha-magma-download' -O {checkpoint_path}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9b4725d",
   "metadata": {},
   "source": [
    "Note that sequence length 1024 checkpoints are also available: magma_v1_1024"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c82287b",
   "metadata": {},
   "outputs": [],
   "source": [
    "session, config, tokenizer = init_inference_session(\n",
    "    \"magma_v1_500\", checkpoint_path=checkpoint_path\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "478b8429",
   "metadata": {},
   "source": [
    "## Run demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee2cd17",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import requests\n",
    "from io import BytesIO\n",
    "import ipywidgets as ipw\n",
    "\n",
    "\n",
    "def answer_int(image_url, text, seed, top_p, top_k, temperature, max_out_tokens):\n",
    "\n",
    "    if image_url.startswith(\"http\"):\n",
    "        response = requests.get(image_url)\n",
    "        image = BytesIO(response.content)\n",
    "    else:\n",
    "        image = open(image_url, \"rb\")\n",
    "    img = ipw.Image(value=image.read(), width=384, height=384)\n",
    "    prompt = ipw.Label(value=f\"Prompt: {text}\", style={\"font_size\": \"16px\"})\n",
    "    answer = ipw.Label(\n",
    "        value=f\"Answer: `{run_inference(session, config, tokenizer, image_url, text, seed, top_p, top_k, temperature, max_out_tokens)}`\",\n",
    "        style={\"font_size\": \"16px\"},\n",
    "    )\n",
    "\n",
    "    return ipw.VBox(\n",
    "        [img, prompt, answer], layout=ipw.Layout(display=\"flex\", align_items=\"center\")\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69f00270",
   "metadata": {},
   "source": [
    "Since the next token selection is done on the CPU, you can change the parameters used for generation (`top_p`, `top_k`, `max_out_tokens`, `temperature`, explained below) without triggering recompilation.\n",
    "Specifying a `seed` allows you to get reproducible deterministic results.\n",
    "\n",
    "Given an image and a text prompt, the model outputs logits for the sequence.\n",
    "Logits corresponding to the last token in the sequence are used to predict the next token.\n",
    "Several heuristics are supported:\n",
    "- top-p: probability is redistributed among the first x tokens such that the cumulative probability is greater than a threshold p. Then, next token is sampled from such distribution (categorical sampling, non deterministic).\n",
    "- top-k: probability is redistributed among the K most-likely tokens. Then, next token is sampled from such distribution (categorical sampling, non deterministic).\n",
    "- temperature: logits are scaled by a factor 1/T (T between 0 and 1 ) before applying the softmax. This makes the distribution more peaked for low temperature, and broader for high temperatures. A zero temperature corresponds to a deterministic choice (argmax), while sampling output becomes more random as we increase the temperature.\n",
    "- max_out_tokens: maximum number of tokens in the answer\n",
    "\n",
    "This [Hugging Face blog post](https://huggingface.co/blog/how-to-generate) provides more information.\n",
    "\n",
    "\n",
    "Manually set model heuristics and the input image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e576d8bc-05a3-4e23-bfcb-2cdad0072ec8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# image_url = \"demo_example_images/cantaloupe_popsicle.jpg\"\n",
    "# image_url=\"demo_example_images/circles.jpg\"\n",
    "# image_url=\"demo_example_images/circles_square.jpg\"\n",
    "# image_url=\"demo_example_images/korea.jpg\"\n",
    "# image_url=\"demo_example_images/matterhorn.jpg\"\n",
    "# image_url=\"demo_example_images/mushroom.jpg\"\n",
    "# image_url=\"demo_example_images/people.jpg\"\n",
    "# image_url=\"demo_example_images/playarea.jpg\"\n",
    "image_url = \"demo_example_images/popsicle.png\"\n",
    "# image_url=\"demo_example_images/rainbow_popsicle.jpeg\"\n",
    "# image_url=\"demo_example_images/table_tennis.jpg\"\n",
    "\n",
    "text = \"A picture of \"\n",
    "seed = 0  # 0 to 300\n",
    "top_p = 0.9  # 0 to 1.0\n",
    "top_k = 0  # 0 to 10\n",
    "temperature = 0.7  # 0 to 1.0\n",
    "max_out_tokens = 6  # 1 to 356\n",
    "\n",
    "answer_int(image_url, text, seed, top_p, top_k, temperature, max_out_tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19d89240",
   "metadata": {},
   "source": [
    "## Next steps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e219de42",
   "metadata": {},
   "source": [
    "Areas for further exploration:\n",
    "* Experiment with the sequence length 1024 checkpoint\n",
    "* Investigate the [GPT-J notebook]() to better understand tensor parallelism on the IPU"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}