notebooks/walkthrough.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3373f2c7-da8e-49eb-8300-25ae4a00fdaa",
   "metadata": {},
   "source": [
    "# Lumigator from Mozilla AI 🐊 🦊"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00ba77ea-aae4-4c55-ad18-1b665836532c",
   "metadata": {},
   "source": [
    "Welcome to the Getting Started notebook for Lumigator, a powerful tool developed by\n",
    "[Mozilla AI](https://www.mozilla.ai/) for evaluating language models. In this guide, we'll\n",
    "walk you through the key features and functionalities of Lumigator, helping you quickly get\n",
    "up to speed with using it for your LM evaluation tasks.\n",
    "\n",
    "Specifically, we'll cover the following topics:\n",
    "\n",
    "+ What Jupyter Notebooks are and how to use them.\n",
    "+ An Overview of Machine Learning.\n",
    "+ The Lumigator Platform 🐊\n",
    "+ How do Machine Learning evaluation workflows look like.\n",
    "+ The Thunderbird Ground Truth Dataset.\n",
    "+ Selecting models to perform summarization:\n",
    "  + Using one encoder/decoder (BART).\n",
    "  + Utilizing two decoders (Mistral and GPT-4) to evaluate against ground truth.\n",
    "+ Running evaluation experiments.\n",
    "+ Discussing results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e41dab1",
   "metadata": {},
   "source": [
    "## Jupyter Walkthrough\n",
    "\n",
    "[Jupyter Notebooks](https://jupyter-notebook.readthedocs.io/en/stable/) provide an executable\n",
    "environment for running (usually) Python code alongside text. To work with Jupyter, click the \"play\"\n",
    "icon (i.e., ▶) to execute the code and view the results below the cell you are currently running.\n",
    "You can also use the `shift + enter` shortcut to execute the code cell and move to the next one.\n",
    "Cells are executed sequentially and can contain either text (Markdown) or Python code.\n",
    "\n",
    "![cell-running](images/running.png)\n",
    "\n",
    "Your files are located on the left-hand side. They'll be saved for the duration of our session, but\n",
    "if you'd like to keep them, make sure to download them. \n",
    "\n",
    "![file-tree](images/files.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4d284df-b55c-43ea-af33-6909392baa0f",
   "metadata": {},
   "source": [
    "Lets' try running some code! Execute the following code and verify the output below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cb7563e-3ab3-486f-b45a-7bd0c47fbcce",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Welcome to Lumigator!🐊\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ecaa846",
   "metadata": {},
   "source": [
    "## Machine learning glossary\n",
    "\n",
    "As we walk through this notebook, you’ll encounter several Machine Learning terms that are essential\n",
    "to understanding the concepts and methods we'll be using. To help you follow along, we've compiled a\n",
    "brief glossary of key terms you'll come across during the session:\n",
    "\n",
    "+ **Machine learning**: The process of creating a model that learns from data.\n",
    "+ **Dataset**: Data used to train models and evaluate their performance.\n",
    "+ **LLM**: Large language model, [a text-based model that performs next-word predictions](https://www.nvidia.com/en-us/glossary/large-language-models/).\n",
    "+ **Tokens**: Words broken up into pieces to be used in an LLM.\n",
    "+ **Inference**: The process of getting a prediction from a large language model.\n",
    "+ **Embeddings**: Numerical representations of text generated by modeling.\n",
    "+ **Encoder-decoder models**: A neural network architecture comprised of two neural networks, an\n",
    "  encoder that takes the input vectors from our data and creates an embedding of a fixed length, and\n",
    "  a decoder, also a neural network, which takes the embeddings encoded as input and generates a\n",
    "  static set of outputs such as translated text or a text summary.\n",
    "+ **Decoder-only models** - Given a fixed input prompt, uses its representation to generate a\n",
    "  sequence of words one at a time, with each word being conditioned on the ones generated\n",
    "  previously.\n",
    "+ **Task** - Machine learning tasks to fit a specific model type, including translation,\n",
    "  summarization, completion, etc.\n",
    "+ **Ground truth** - Information that has been evaluated to be true by humans (or LLMs, in some\n",
    "  cases), that we can use to evaluate and compare trained models.\n",
    "\n",
    "The process of **machine learning** is the process of creating a mathematical model that tries to\n",
    "approximate the world. A **machine learning model** is a set of instructions for generating a given\n",
    "output from data. The instructions are learned from the features of the input data itself.\n",
    "\n",
    "Within the broad landscape of machine learning, there are various modeling approaches, including\n",
    "**supervised**, **unsupervised**, and **reinforcement** learning. Each approach has its own set of\n",
    "techniques and use cases.\n",
    "\n",
    "In the context of language models, the kind of modeling we focus on with Large Language Models\n",
    "(LLMs) primarily falls within the domain of neural network-based approaches. These models learn\n",
    "patterns, structures, and relationships in data through vast networks of interconnected nodes (or\n",
    "neurons), which allow them to generate, interpret, and manipulate natural language in powerful ways."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f67ff2a",
   "metadata": {},
   "source": [
    "### How do LLMs work? \n",
    "\n",
    "There are many different kinds of LLMs and many different kinds of architectures. For our\n",
    "evaluations, we use two different kinds:\n",
    "\n",
    "+ **Encoder/Decoder** - BART is an encoder/decoder model that converts input data into a fixed-size\n",
    "  representation (similar to encoder models). These models are trained first to transform text into\n",
    "  numerical representations, then to output text based on those numerical representations. They're\n",
    "  good for synthesis as opposed to generation. \n",
    "+ **Decoder-only** - most models in the GPT-family, like Mistral, GPT, and others we'll be working\n",
    "  with, are pre-trained with text data in an autoregressive manner, for next-token prediction given\n",
    "  previous tokens.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8139990",
   "metadata": {},
   "source": [
    "### LLM Evaluation Workflows\n",
    "\n",
    "The following steps outline the key phases involved in evaluating a Large Language Model (LLM):\n",
    "\n",
    "1. **Generate Ground Truth**: The first step is to establish a reliable ground truth based on the\n",
    "   specific business use case you're targeting. This represents the \"correct\" or expected output\n",
    "   against which model performance will be measured.\n",
    "1. **Select Models for Evaluation**: Next, choose several candidate models that you'd like to\n",
    "   evaluate. These could be different versions of a language model, or distinct models altogether,\n",
    "   depending on your evaluation criteria and use case.\n",
    "1. **Run the Evaluation Loop**: This phase involves running the models through an evaluation process\n",
    "   where you compare the model's outputs against the ground truth. You'll iterate through multiple\n",
    "   examples, assessing how well each model performs in generating the desired results.\n",
    "1. **Analyze Evaluation Results**: Finally, after completing the evaluation loop, you'll analyze the\n",
    "   results to identify strengths, weaknesses, and areas for improvement. This analysis helps inform\n",
    "   decisions about which model to use, optimize, or further train for the business use case at hand.\n",
    "\n",
    "![lumigator-features](images/lumigator-features.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c44d8a1",
   "metadata": {},
   "source": [
    "On a technical level, Lumigator is a Python-based FastAPI web application designed to run services\n",
    "that handle jobs and deployments on a Ray cluster. This cluster can be run either locally or in the\n",
    "cloud, depending on your system's specifications and the resources available.\n",
    "\n",
    "The results and job metadata generated during evaluations are stored in an SQL database for easy\n",
    "tracking and retrieval. Additionally, larger models, which are often loaded from platforms like\n",
    "[Hugging Face](https://huggingface.co/), require GPUs for efficient processing due to their size and\n",
    "computational demands.\n",
    "\n",
    "![lumigator-architecture](images/lumigator-architecture.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ed38e49",
   "metadata": {},
   "source": [
    "What is Ray? [A distributed runtime for Python programs](https://github.com/ray-project/ray) that\n",
    "includes a Core library with primitives (Tasks, Actors, and objects) and a suite of ML libraries\n",
    "(Tune, Serve) that allow to build components of the Machine Learning model workflow. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "838d2fd4",
   "metadata": {},
   "source": [
    "### Nota bene: Machine Learning is alchemy\n",
    "\n",
    "When we think of traditional software application workflows, we think of an example such as adding a\n",
    "button. We can clearly test that we've added a blue button to our application, and that it works\n",
    "correctly. Machine Learning is not like this! It involves a lot of experimentation, tweaking of\n",
    "hyperparameters and prompts and trying different models. Expect for the process to be imperfect,\n",
    "with many iterative loops. Luckily, Lumigator helps take away the uncertainty of at least model\n",
    "selection. 🙂\n",
    "\n",
    "> There’s a self-congratulatory feeling in the air. We say things like “Machine Learning is the new\n",
    "> electricity”. I’d like to offer an alternative metaphor: Machine Learning has become alchemy. -\n",
    "> [Ben Recht and Ali Rahimi](https://archives.argmin.net/2017/12/05/kitchen-sinks/)\n",
    "\n",
    "Ultimately, the final conclusion of whether a model is good is if humans think it's good. With that\n",
    "in mind, let's dive into setting up experiments with Lumigator to test our models!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd36c70c-ac15-4d2f-a890-ac84c6a374bc",
   "metadata": {},
   "source": [
    "## Importing Required Libraries\n",
    "\n",
    "Before we begin working with Lumigator, we'll need to import several libraries and modules that\n",
    "provide the necessary functionality for our tasks. In this section, we'll load the core\n",
    "dependencies, including tools for data handling, and visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f7afa9ec9b5f866",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing packages we need to work with data\n",
    "# python standard libraries\n",
    "import os\n",
    "\n",
    "# Random string generator\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# third-party libraries\n",
    "import pandas as pd\n",
    "from datasets import load_dataset\n",
    "from lumigator_schemas.datasets import DatasetFormat\n",
    "from lumigator_sdk.lumigator import LumigatorClient\n",
    "from lumigator_sdk.strict_schemas import JobCreate, JobEvalConfig, JobInferenceConfig\n",
    "from utils import get_nested_value, job_result_download, results_to_table\n",
    "\n",
    "# wrap columns for inspection\n",
    "pd.set_option(\"display.max_colwidth\", 0)\n",
    "# stylesheet for visibility\n",
    "plt.style.use(\"fast\")\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50d76aec",
   "metadata": {},
   "outputs": [],
   "source": [
    "LUMIGATOR_SERVICE_HOST = os.getenv(\"LUMIGATOR_SERVICE_HOST\", \"localhost\")\n",
    "LUMIGATOR_SERVICE_PORT = os.getenv(\"LUMIGATOR_SERVICE_PORT\", \"8000\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f1cd443f885af1d",
   "metadata": {},
   "source": [
    "# Understanding the Lumigator App and API \n",
    "\n",
    "The app itself consists of an API, which you can access and test out methods in the\n",
    "[OpenAPI spec](https://swagger.io/specification/), at the platform URL, under docs.\n",
    "\n",
    "If you are running Lumigator as a local installation, you can directly access the API at\n",
    "[this URL](http://localhost:8000/docs).\n",
    "\n",
    "![lumigator-api](images/lumigator-api.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd6f7aa716833f7",
   "metadata": {},
   "source": [
    "Large language models today are consumed in one of several ways:\n",
    "\n",
    "+ As **API endpoints** for proprietary models hosted by [OpenAI](https://openai.com/),\n",
    "  [Anthropic](https://www.anthropic.com/), or major cloud providers.\n",
    "+ As **model artifacts** downloaded from HuggingFace’s Model Hub, trained/fine-tuned using\n",
    "  HuggingFace libraries, and hosted on local storage.\n",
    "+ As model artifacts available in a format optimized for **local inference**, typically\n",
    "  [GGUF](https://github.com/ggerganov/ggml/blob/master/docs/gguf.md), and accessed via applications\n",
    "  like [llama.cpp](https://github.com/ggerganov/llama.cpp) or [ollama](https://ollama.com/).\n",
    "+ As [ONNX](https://onnx.ai/), a format which optimizes sharing between backend ML frameworks.\n",
    "\n",
    "We use API endpoints and local storage in Lumigator. We currently have four key endpoints on the\n",
    "platform:\n",
    "\n",
    "+ `/health`: Status of the application, running status of jobs and deployments. \n",
    "+ `/datasets`: Data that we add to the platform for evaluation. We can upload, delete, and save\n",
    "  different data in the platform. We'll use this to also save our ground truth and experiment data.\n",
    "+ `/jobs`: Our actual jobs. We can list all previous jobs, create new ones, and\n",
    "  get their results. This includes evaluation as well as inference access to external APIs such as Mistral and OpenAI."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1a28c048ea00706",
   "metadata": {},
   "source": [
    "## Model Task: Summarization\n",
    "\n",
    "The task we'll be working with is *summarization*, aka we want to generate a summary of our text. In\n",
    "our business case, which is to create summaries of conversation threads, much as you might see in\n",
    "Slack or an email chain, the models need to be able to extract key information from those threads\n",
    "while still being able to accept a large context window to capture the entire conversation history. \n",
    "\n",
    "We identified that it is far more valuable to conduct *abstractive* summaries—summaries that\n",
    "identify important sections in the text and generate highlights—rather than *extractive* ones, which\n",
    "select a subset of sentences and staple them together. This is because the final interface will be\n",
    "in natural language, and we want to avoid summaries that are interpreted from often incoherent text\n",
    "snippets produced by extractive methods.\n",
    "\n",
    "For more on summarization as a use-case, [see our blog post here.](https://blog.mozilla.ai/on-model-selection-for-text-summarization/)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f33e80b64bad99e0",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Ground Truth for Models\n",
    "\n",
    "The term ground truth comes from geology and geospatial sciences, where actual information was\n",
    "collected on the ground to validate data acquired through remote sensing, such as satellite imagery\n",
    "or aerial photography. Since then, the concept has been adopted in other fields, particularly in\n",
    "machine learning and artificial intelligence, to refer to the accurate, real-world data used for\n",
    "training and testing models. \n",
    "\n",
    "The **best ground truth is human-generated** but building it is a very expensive task. One recent\n",
    "trend is to rely on large language models but they have their own pitfalls. An intermediate approach\n",
    "uses different LLMs to provide ground truth \"candidates\" which are then subject to human pairwise\n",
    "evaluation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86d102a0b4515342",
   "metadata": {},
   "source": [
    "## Our Input data\n",
    "\n",
    "The data we'll be using in this walkthrough comes from\n",
    "[DialogSum](https://github.com/cylnlp/DialogSum), a large-scale labeled dialogue summarization\n",
    "dataset which comes with ground truth provided by human annotators. Here follows a brief description\n",
    "of DialogSum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65af6945-bd5d-4506-af3c-144683779d58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The dataset is available at https://huggingface.co/datasets/knkarthick/dialogsum\n",
    "# and can be directly downloaded with the `load_dataset` method\n",
    "dataset = \"knkarthick/dialogsum\"\n",
    "ds = load_dataset(dataset, split=\"validation\")\n",
    "df = ds.to_pandas()  # noqa: PD901"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cb05409-aad2-42e2-bc23-743575ad924e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Examine a single sample\n",
    "df[\"dialogue\"].iloc[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac02084ba6b6fd5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a function to do some simple character counts for model input\n",
    "df[\"char_count\"] = df[\"dialogue\"].str.len()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c90ceb9398ae3ed0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# inspect our data\n",
    "df.head(n=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "359677f7f12bcefd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show statistics about characters count\n",
    "df[\"char_count\"].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "116e1127c3846883",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate plot of character counts\n",
    "fig, ax = plt.subplots(figsize=(12, 6))\n",
    "ax.hist(df[\"char_count\"], bins=30)\n",
    "ax.set_xlabel(\"Character Count\")\n",
    "ax.set_ylabel(\"Frequency\")\n",
    "\n",
    "stats = df[\"char_count\"].describe().apply(lambda x: f\"{x:.0f}\")\n",
    "\n",
    "# Add text boxes for statistics\n",
    "plt.text(1.05, 0.95, stats.to_string(), transform=ax.transAxes, verticalalignment=\"top\")\n",
    "\n",
    "# Adjust layout\n",
    "plt.tight_layout()\n",
    "fig.subplots_adjust(right=0.75)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "378251f3-4145-4fa6-b933-564a91101c44",
   "metadata": {},
   "source": [
    "## Save and upload datasets\n",
    "\n",
    "Now that you have seen how the option to generate ground truth works, let us save all datasets and\n",
    "make them available to lumigator for further experiments. For each example (i.e., dialogsum original\n",
    "dataset) we will perform the following operations:\n",
    "\n",
    "1. Make sure that the two main fields (original text and ground truth) are called `examples` and\n",
    "  `ground_truth`, which are the names internally used by Lumigator to refer to them, and save the\n",
    "  datasets as CSV files.\n",
    "2. Make the dataset available to Lumigator with the `create_dataset` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e0f300c-9651-473b-aca9-1f147a3f2f2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = ds.remove_columns([\"id\", \"topic\"])\n",
    "ds = ds.rename_column(\"dialogue\", \"examples\")\n",
    "ds = ds.rename_column(\"summary\", \"ground_truth\")\n",
    "\n",
    "dataset_name = \"dialogsum_converted.csv\"\n",
    "ds.to_csv(dataset_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3ec30ff-f639-4ab4-b97d-eadbe3df8075",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from urllib3 import Retry\n",
    "\n",
    "notebok_retry = Retry(\n",
    "    connect=10,\n",
    "    backoff_factor=1,\n",
    "    backoff_max=60,\n",
    ")\n",
    "\n",
    "lm_client = LumigatorClient(f\"{LUMIGATOR_SERVICE_HOST}:{LUMIGATOR_SERVICE_PORT}\", retry_conf=notebok_retry)\n",
    "\n",
    "loaded_dataset = lm_client.datasets.create_dataset(Path.open(dataset_name, \"rb\"), DatasetFormat.JOB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26e5b2701933bc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# And let's check that data loaded\n",
    "datasets = lm_client.datasets.get_datasets()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a152e1e084e7b64",
   "metadata": {},
   "source": [
    "## Jobs\n",
    "\n",
    "After generating the ground truth (either manually or with the aid of some models) and uploading the\n",
    "dataset to lumigator, we are ready to start evaluating models on it. Note that when you uploaded\n",
    "your datasets you got back some information that included a dataset `id`. This is a unique\n",
    "identifier to your own dataset that you can reuse across different jobs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c98f8a6fdff48654",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = loaded_dataset\n",
    "\n",
    "# now look for the dataset on lumigator\n",
    "result = lm_client.datasets.get_dataset(dataset.id)\n",
    "dataset_id, dataset_name = result.id, result.filename\n",
    "print(f\"Annotated dataset: {dataset_id, dataset_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a51eb04dbc16b035",
   "metadata": {},
   "source": [
    "### Model Selection\n",
    "\n",
    "What you see below are different lists of models we have already tested for the summarization task.\n",
    "The `models` variable at the end provides you with a selection, but you can choose any combination\n",
    "of them: the default is a single local model (`facebook/bart-large-cnn`), but depending on your\n",
    "setup you can choose more and/or add different APIs.\n",
    "\n",
    "Note that different model types are specified with different prefixes:\n",
    "\n",
    "- `hf` is used for HuggingFace models which are downloaded and ran as Ray jobs.\n",
    "- `mistral/` is used for models which are accessed through the Mistral API.\n",
    "- `gpt-4` is specified directly per the [LiteLLM api](https://docs.litellm.ai/docs/providers/openai)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af3f8c76c37a2b0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here follows a list of models we have tested for summarization:\n",
    "# feel free to add any of them in the \"models\" list below\n",
    "#\n",
    "# Encoder-Decoder models\n",
    "#    'facebook/bart-large-cnn',\n",
    "#    'Falconsai/text_summarization',\n",
    "#\n",
    "# Decoder models\n",
    "#    'mistral/open-mistral-7b',\n",
    "#\n",
    "# GPTs\n",
    "#    \"gpt-4o-mini\",\n",
    "#    \"gpt-4-turbo\",\n",
    "#    \"gpt-3.5-turbo-0125\",\n",
    "#\n",
    "\n",
    "# We'll use this one just for testing quickly, the results will be random\n",
    "models = [\n",
    "    \"hf-internal-testing/tiny-random-BARTForConditionalGeneration\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fa55eee63016041",
   "metadata": {},
   "source": [
    "### Run Evaluations\n",
    "\n",
    "The following cell will start the actual model evaluations. Once you run it, new jobs will be\n",
    "submitted to ray (one for each model) and the outcomes of these submissions will be printed.\n",
    "\n",
    "Each evaluation job will first use the provided model to summarize each of the emails in the input\n",
    "dataset. After that, it will calculate a few metrics to evaluate how close the predicted summaries\n",
    "are to the ground truth provided in the dataset. Each job starts with a `created` status. While the\n",
    "job runs, you will be able to track its status by running the cell in the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f344b6c4998e409",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set this value to limit the evaluation to the first max_samples items (0=all)\n",
    "max_samples = 10\n",
    "# team_name is a way to group jobs together under the same namespace, feel free to customize it\n",
    "team_name = \"lumigator_enthusiasts\"\n",
    "\n",
    "responses = []\n",
    "for model in models:\n",
    "    infer_job_config = JobInferenceConfig(\n",
    "        model=model,\n",
    "        provider=\"hf\",\n",
    "        output_field=\"predictions\",\n",
    "        store_to_dataset=True,\n",
    "    )\n",
    "    infer_job_create = JobCreate(\n",
    "        name=f\"{team_name}: Summary inference for model {model}\",\n",
    "        description=\"Test run (inference)\",\n",
    "        dataset=dataset_id,\n",
    "        max_samples=max_samples,\n",
    "        job_config=infer_job_config,\n",
    "    )\n",
    "    job_infer_creation_result = lm_client.jobs.create_job(infer_job_create)\n",
    "    lm_client.jobs.wait_for_job(job_infer_creation_result.id)\n",
    "\n",
    "    infer_dataset = lm_client.jobs.get_job_dataset(str(job_infer_creation_result.id))\n",
    "\n",
    "    eval_job_config = JobEvalConfig(\n",
    "        metrics=[\"rouge\", \"meteor\", \"bertscore\"],\n",
    "    )\n",
    "    eval_job_create = JobCreate(\n",
    "        name=f\"{team_name}: Summary evaluation for model {model}\",\n",
    "        description=\"Test run (evaluation)\",\n",
    "        dataset=str(infer_dataset.id),\n",
    "        job_config=eval_job_config,\n",
    "    )\n",
    "\n",
    "    job_eval_creation_result = lm_client.jobs.create_job(eval_job_create)\n",
    "    responses.append(job_eval_creation_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a250dc71-daf7-4644-a8d5-e345bbb3c61a",
   "metadata": {},
   "source": [
    "![ray-job](images/ray-job.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f787190-4581-4c26-beca-0260303a68bd",
   "metadata": {},
   "source": [
    "### Track Evaluation Jobs\n",
    "\n",
    "To track the progress of your evaluation jobs, you’ll need to run the following commands. These jobs\n",
    "are executed on a Ray cluster, which efficiently distributes the workload across multiple nodes,\n",
    "whether locally or in the cloud.\n",
    "\n",
    "Not that you won't be able to run other cells while this one is running. However, you can interrupt\n",
    "it whenever you want by clicking on the \"stop\" button above and run it at a later time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2858bcea-6fcb-4441-b56e-5d89502f10e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for response in responses:\n",
    "    job_id = response.id\n",
    "    job = lm_client.jobs.wait_for_job(job_id)\n",
    "    print(job)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32d8cbf3-e44f-4f1e-b520-3f6fac147fca",
   "metadata": {},
   "source": [
    "## Show evaluation results\n",
    "\n",
    "Once all evaluations are completed, their results will be stored on our platform and available for\n",
    "download. The following cell iterates on all your job ids, downloads results from each, and builds a\n",
    "table comparing different metrics for each model.\n",
    "\n",
    "The metrics we use to evaluate are ROUGE, METEOR, and BERT score. They all measure similarity\n",
    "between predicted summaries and those provided with the ground truth, but each of them focuses on\n",
    "different aspects. The image below shows their main characteristics and the tradeoffs between their\n",
    "flexibility and their computational cost.\n",
    "\n",
    "![metrics](images/metrics.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6163cf54-0719-4194-abc3-fb03d0b8ec37",
   "metadata": {},
   "outputs": [],
   "source": [
    "job_id = responses[0].id\n",
    "job_download = lm_client.jobs.get_job_download(job_id)\n",
    "result = job_result_download(job_download)\n",
    "print(result)\n",
    "model_name = models[0]\n",
    "result[\"artifacts\"][\"model\"] = model_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7a828ad-089b-4a2c-95f1-a00d1c34a366",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_to_table([(result, model_name)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17b450ce",
   "metadata": {},
   "source": [
    "## Analysis of Evaluation Results\n",
    "\n",
    "The table above is just a summary of all the evaluation results. The `result` object contains way\n",
    "more details from which you'll be able to get a few more insights in the following cells.\n",
    "\n",
    "The following cell shows you the kind of information that's available in each of the `result` object\n",
    "elements. This information is nested at different depth levels. You can access each using the\n",
    "`get_nested_value` command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eadbdaec-a1e9-4731-bfdd-b471e6a0fbcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# eval_results is a list holding information for each of the models you defined before\n",
    "# for each element, you can access different metrics, time performance, and predictions\n",
    "result.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ec4d37e-4750-404e-b40c-3778b8196d64",
   "metadata": {},
   "outputs": [],
   "source": [
    "# see how much time it took for a model to summarize all the input samples\n",
    "get_nested_value(result, \"inference_time\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13a73261-e5a4-46a6-9b7c-1bc5e62dc1ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# see all the bertscore data\n",
    "get_nested_value(result, \"bertscore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f6e8285-f407-4a5f-acda-3400554c3d2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# see mean bert precision\n",
    "get_nested_value(result, \"bertscore/precision_mean\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}