{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implement custom tests\n",
    "\n",
    "Custom tests extend the functionality of ValidMind, allowing you to document any model or use case with added flexibility.\n",
    "\n",
    "ValidMind provides a comprehensive set of tests out-of-the-box to evaluate and document your models and datasets. We recognize there will be cases where the default tests do not support a model or dataset, or specific documentation is needed. In these cases, you can create and use your own custom code to accomplish what you need. To streamline custom code integration, we support the creation of custom test functions.\n",
    "\n",
    "This interactive notebook provides a step-by-step guide for implementing and registering custom tests with ValidMind, running them individually, viewing the results on the ValidMind Platform, and incorporating them into your model documentation template.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents \n",
    "- [About ValidMind](#toc2_)    \n",
    "  - [Before you begin](#toc2_1_)    \n",
    "  - [New to ValidMind?](#toc2_2_)    \n",
    "  - [Key concepts](#toc2_3_)    \n",
    "- [Before you begin](#toc3_)    \n",
    "  - [New to ValidMind?](#toc3_1_)    \n",
    "- [Install the ValidMind Library](#toc4_)    \n",
    "- [Initialize the ValidMind Library](#toc5_)    \n",
    "  - [Get your code snippet](#toc5_1_) \n",
    "- [Implement a Custom Test](#toc6_)    \n",
    "- [Run the Custom Test](#toc7_)    \n",
    "  - [Setup the Model and Dataset](#toc7_1_)    \n",
    "  - [Run the Custom Test](#toc7_2_)    \n",
    "- [Adding Custom Tests to Model Documentation](#toc8_)    \n",
    "- [Some More Custom Tests](#toc9_)    \n",
    "  - [Custom Test: Table of Model Hyperparameters](#toc9_1_)    \n",
    "  - [Custom Test: External API Call](#toc9_2_)    \n",
    "  - [Custom Test: Passing Parameters](#toc9_3_)    \n",
    "  - [Custom Test: Multiple Tables and Plots in a Single Test](#toc9_4_)    \n",
    "  - [Custom Test: Images](#toc9_5_)\n",
    "  - [Custom Test: Description](#toc9_6_)\n",
    "\n",
    "- [Conclusion](#toc10_)    \n",
    "- [Next steps](#toc11_)    \n",
    "  - [Work with your model documentation](#toc11_1_)    \n",
    "  - [Discover more learning resources](#toc11_2_)    \n",
    "- [Upgrade ValidMind](#toc12_)\n",
    "\n",
    "<!-- vscode-jupyter-toc-config\n",
    "\tnumbering=false\n",
    "\tanchor=true\n",
    "\tflat=false\n",
    "\tminLevel=2\n",
    "\tmaxLevel=4\n",
    "\t/vscode-jupyter-toc-config -->\n",
    "<!-- THIS CELL WILL BE REPLACED ON TOC UPDATE. DO NOT WRITE YOUR TEXT IN THIS CELL -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc2_'></a>\n",
    "\n",
    "## About ValidMind\n",
    "\n",
    "ValidMind is a suite of tools for managing model risk, including risk associated with AI and statistical models.\n",
    "\n",
    "You use the ValidMind Library to automate documentation and validation tests, and then use the ValidMind Platform to collaborate on model documentation. Together, these products simplify model risk management, facilitate compliance with regulations and institutional standards, and enhance collaboration between yourself and model validators.\n",
    "\n",
    "<a id='toc2_1_'></a>\n",
    "\n",
    "### Before you begin\n",
    "\n",
    "This notebook assumes you have basic familiarity with Python, including an understanding of how functions work. If you are new to Python, you can still run the notebook but we recommend further familiarizing yourself with the language. \n",
    "\n",
    "If you encounter errors due to missing modules in your Python environment, install the modules with `pip install`, and then re-run the notebook. For more help, refer to [Installing Python Modules](https://docs.python.org/3/installing/index.html).\n",
    "\n",
    "<a id='toc2_2_'></a>\n",
    "\n",
    "### New to ValidMind?\n",
    "\n",
    "If you haven't already seen our documentation on the [ValidMind Library](https://docs.validmind.ai/developer/validmind-library.html), we recommend you begin by exploring the available resources in this section. There, you can learn more about documenting models and running tests, as well as find code samples and our Python Library API reference.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\" style=\"background-color: #B5B5B510; color: black; border: 1px solid #083E44; border-left-width: 5px; box-shadow: 2px 2px 4px rgba(0, 0, 0, 0.2);border-radius: 5px;\"><span style=\"color: #083E44;\"><b>For access to all features available in this notebook, create a free ValidMind account.</b></span>\n",
    "<br></br>\n",
    "Signing up is FREE — <a href=\"https://docs.validmind.ai/guide/configuration/register-with-validmind.html\" style=\"color: #DE257E;\"><b>Register with ValidMind</b></a></div>\n",
    "\n",
    "<a id='toc2_3_'></a>\n",
    "\n",
    "### Key concepts\n",
    "\n",
    "**Model documentation**: A structured and detailed record pertaining to a model, encompassing key components such as its underlying assumptions, methodologies, data sources, inputs, performance metrics, evaluations, limitations, and intended uses. It serves to ensure transparency, adherence to regulatory requirements, and a clear understanding of potential risks associated with the model’s application.\n",
    "\n",
    "**Documentation template**: Functions as a test suite and lays out the structure of model documentation, segmented into various sections and sub-sections. Documentation templates define the structure of your model documentation, specifying the tests that should be run, and how the results should be displayed.\n",
    "\n",
    "**Tests**: A function contained in the ValidMind Library, designed to run a specific quantitative test on the dataset or model. Tests are the building blocks of ValidMind, used to evaluate and document models and datasets, and can be run individually or as part of a suite defined by your model documentation template.\n",
    "\n",
    "**Custom tests**: Custom tests are functions that you define to evaluate your model or dataset. These functions can be registered via the ValidMind Library to be used with the ValidMind Platform.\n",
    "\n",
    "**Inputs**: Objects to be evaluated and documented in the ValidMind Library. They can be any of the following:\n",
    "\n",
    "  - **model**: A single model that has been initialized in ValidMind with [`vm.init_model()`](https://docs.validmind.ai/validmind/validmind.html#init_model).\n",
    "  - **dataset**: Single dataset that has been initialized in ValidMind with [`vm.init_dataset()`](https://docs.validmind.ai/validmind/validmind.html#init_dataset).\n",
    "  - **models**: A list of ValidMind models - usually this is used when you want to compare multiple models in your custom test.\n",
    "  - **datasets**: A list of ValidMind datasets - usually this is used when you want to compare multiple datasets in your custom test. See this [example](https://docs.validmind.ai/notebooks/how_to/run_tests_that_require_multiple_datasets.html) for more information.\n",
    "\n",
    "**Parameters**: Additional arguments that can be passed when running a ValidMind test, used to pass additional information to a test, customize its behavior, or provide additional context.\n",
    "\n",
    "**Outputs**: Custom test can return elements like tables or plots. Tables may be a list of dictionaries (each representing a row) or a pandas DataFrame. Plots may be matplotlib or plotly figures.\n",
    "\n",
    "**Test suites**: Collections of tests designed to run together to automate and generate model documentation end-to-end for specific use-cases.\n",
    "\n",
    "Example: the [`classifier_full_suite`](https://docs.validmind.ai/validmind/validmind/test_suites/classifier.html#ClassifierFullSuite) test suite runs tests from the [`tabular_dataset`](https://docs.validmind.ai/validmind/validmind/test_suites/tabular_datasets.html) and [`classifier`](https://docs.validmind.ai/validmind/validmind/test_suites/classifier.html) test suites to fully document the data and model sections for binary classification model use-cases.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc4_'></a>\n",
    "\n",
    "## Install the ValidMind Library\n",
    "\n",
    "To install the library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q validmind"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc5_'></a>\n",
    "\n",
    "## Initialize the ValidMind Library\n",
    "\n",
    "ValidMind generates a unique _code snippet_ for each registered model to connect with your developer environment. You initialize the ValidMind Library with this code snippet, which ensures that your documentation and tests are uploaded to the correct model when you run the notebook.\n",
    "\n",
    "<a id='toc5_1_'></a>\n",
    "\n",
    "### Get your code snippet\n",
    "\n",
    "1. In a browser, [log in to ValidMind](https://docs.validmind.ai/guide/configuration/log-in-to-validmind.html).\n",
    "\n",
    "2. In the left sidebar, navigate to **Model Inventory** and click **+ Register Model**.\n",
    "\n",
    "3. Enter the model details and click **Continue**. ([Need more help?](https://docs.validmind.ai/guide/model-inventory/register-models-in-inventory.html))\n",
    "\n",
    "   For example, to register a model for use with this notebook, select:\n",
    "\n",
    "   - Documentation template: `Binary classification`\n",
    "   - Use case: `Marketing/Sales - Attrition/Churn Management`\n",
    "\n",
    "   You can fill in other options according to your preference.\n",
    "\n",
    "4. Go to **Getting Started** and click **Copy snippet to clipboard**.\n",
    "\n",
    "Next, [load your model identifier credentials from an `.env` file](https://docs.validmind.ai/developer/model-documentation/store-credentials-in-env-file.html) or replace the placeholder with your own code snippet:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load your model identifier credentials from an `.env` file\n",
    "\n",
    "%load_ext dotenv\n",
    "%dotenv .env\n",
    "\n",
    "# Or replace with your code snippet\n",
    "\n",
    "import validmind as vm\n",
    "\n",
    "vm.init(\n",
    "    # api_host=\"...\",\n",
    "    # api_key=\"...\",\n",
    "    # api_secret=\"...\",\n",
    "    # model=\"...\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc6_'></a>\n",
    "\n",
    "## Implement a Custom Test\n",
    "\n",
    "Let's start off by creating a simple custom test that creates a Confusion Matrix for a binary classification model. We will use the `sklearn.metrics.confusion_matrix` function to calculate the confusion matrix and then display it as a heatmap using `plotly`. (This is already a built-in test in ValidMind, but we will use it as an example to demonstrate how to create custom tests.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from sklearn import metrics\n",
    "\n",
    "\n",
    "@vm.test(\"my_custom_tests.ConfusionMatrix\")\n",
    "def confusion_matrix(dataset, model):\n",
    "    \"\"\"The confusion matrix is a table that is often used to describe the performance of a classification model on a set of data for which the true values are known.\n",
    "\n",
    "    The confusion matrix is a 2x2 table that contains 4 values:\n",
    "\n",
    "    - True Positive (TP): the number of correct positive predictions\n",
    "    - True Negative (TN): the number of correct negative predictions\n",
    "    - False Positive (FP): the number of incorrect positive predictions\n",
    "    - False Negative (FN): the number of incorrect negative predictions\n",
    "\n",
    "    The confusion matrix can be used to assess the holistic performance of a classification model by showing the accuracy, precision, recall, and F1 score of the model on a single figure.\n",
    "    \"\"\"\n",
    "    y_true = dataset.y\n",
    "    y_pred = dataset.y_pred(model)\n",
    "\n",
    "    confusion_matrix = metrics.confusion_matrix(y_true, y_pred)\n",
    "\n",
    "    cm_display = metrics.ConfusionMatrixDisplay(\n",
    "        confusion_matrix=confusion_matrix, display_labels=[False, True]\n",
    "    )\n",
    "    cm_display.plot()\n",
    "\n",
    "    plt.close()  # close the plot to avoid displaying it\n",
    "\n",
    "    return cm_display.figure_  # return the figure object itself"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thats our custom test defined and ready to go... Let's take a look at whats going on here:\n",
    "\n",
    "- The function `confusion_matrix` takes two arguments `dataset` and `model`. This is a VMDataset and VMModel object respectively.\n",
    "- The function docstring provides a description of what the test does. This will be displayed along with the result in this notebook as well as in the ValidMind Platform.\n",
    "- The function body calculates the confusion matrix using the `sklearn.metrics.confusion_matrix` function and then plots it using `sklearn.metric.ConfusionMatrixDisplay`.\n",
    "- The function then returns the `ConfusionMatrixDisplay.figure_` object - this is important as the ValidMind Library expects the output of the custom test to be a plot or a table.\n",
    "- The `@vm.test` decorator is doing the work of creating a wrapper around the function that will allow it to be run by the ValidMind Library. It also registers the test so it can be found by the ID `my_custom_tests.ConfusionMatrix` (see the section below on how test IDs work in ValidMind and why this format is important)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc7_'></a>\n",
    "\n",
    "## Run the Custom Test\n",
    "\n",
    "Now that we have defined and registered our custom test, lets see how we can run it and properly use it in the ValidMind Platform."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc7_1_'></a>\n",
    "\n",
    "### Setup the Model and Dataset\n",
    "\n",
    "First let's setup a an example model and dataset to run our custom metic against. Since this is a Confusion Matrix, we will use the Customer Churn dataset that ValidMind provides and train a simple XGBoost model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xgboost as xgb\n",
    "from validmind.datasets.classification import customer_churn\n",
    "\n",
    "raw_df = customer_churn.load_data()\n",
    "train_df, validation_df, test_df = customer_churn.preprocess(raw_df)\n",
    "\n",
    "x_train = train_df.drop(customer_churn.target_column, axis=1)\n",
    "y_train = train_df[customer_churn.target_column]\n",
    "x_val = validation_df.drop(customer_churn.target_column, axis=1)\n",
    "y_val = validation_df[customer_churn.target_column]\n",
    "\n",
    "model = xgb.XGBClassifier(early_stopping_rounds=10)\n",
    "model.set_params(\n",
    "    eval_metric=[\"error\", \"logloss\", \"auc\"],\n",
    ")\n",
    "model.fit(\n",
    "    x_train,\n",
    "    y_train,\n",
    "    eval_set=[(x_val, y_val)],\n",
    "    verbose=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Easy enough! Now we have a model and dataset setup and trained. One last thing to do is bring the dataset and model into the ValidMind Library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for now, we'll just use the test dataset\n",
    "vm_test_ds = vm.init_dataset(\n",
    "    dataset=test_df,\n",
    "    target_column=customer_churn.target_column,\n",
    "    input_id=\"test_dataset\",\n",
    ")\n",
    "\n",
    "vm_model = vm.init_model(model, input_id=\"model\")\n",
    "\n",
    "# link the model to the dataset\n",
    "vm_test_ds.assign_predictions(model=vm_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc7_2_'></a>\n",
    "\n",
    "### Run the Custom Test\n",
    "\n",
    "Now that we have our model and dataset setup, we have everything we need to run our custom test. We can do this by importing the `run_test` function from the `validmind.tests` module and passing in the test ID of our custom test along with the model and dataset we want to run it against.\n",
    "\n",
    ">Notice how the `inputs` dictionary is used to map an `input_id` which we set above to the `model` and `dataset` keys that are expected by our custom test function. This is how the ValidMind Library knows which inputs to pass to different tests and is key when using many different datasets and models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from validmind.tests import run_test\n",
    "\n",
    "result = run_test(\n",
    "    \"my_custom_tests.ConfusionMatrix\",\n",
    "    inputs={\"model\": \"model\", \"dataset\": \"test_dataset\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You'll notice that the docstring becomes a markdown description of the test. The figure is then displayed as the test result. What you see above is how it will look in the ValidMind Platform as well. Let's go ahead and log the result to see how that works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.log()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc8_'></a>\n",
    "\n",
    "## Adding Custom Test to Model Documentation\n",
    "\n",
    "To do this, go to the documentation page of the model you registered above and navigate to the `Model Development` -> `Model Evaluation` section. Then hover between any existing content block to reveal the `+` button as shown in the screenshot below.\n",
    "\n",
    "![screenshot showing insert button for test-driven blocks](../../images/insert-test-driven-block.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now click on the `+` button and select the `Test-Driven Block` option. This will open a dialog where you can select `My Custom Tests Confusion Matrix` from the list of available tests. You can preview the result and then click `Insert Block` to add it to the documentation.\n",
    "\n",
    "![screenshot showing how to insert a test-driven block](../../images/insert-test-driven-block-custom.png)\n",
    "\n",
    "The test should match the result you see above. It is now part of your documentation and will now be run everytime you run `vm.run_documentation_tests()` for your model. Let's do that now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "vm.reload()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you preview the template, it should show the custom test in the `Model Development`->`Model Evaluation` section:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vm.preview_template()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just so we can run all of the tests in the template, let's initialize the train and raw dataset.\n",
    "\n",
    "(Refer to [**Quickstart for model documentation**](../../quickstart/quickstart_model_documentation.ipynb) and the ValidMind docs for more information on what we are doing here)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vm_raw_dataset = vm.init_dataset(\n",
    "    dataset=raw_df,\n",
    "    input_id=\"raw_dataset\",\n",
    "    target_column=customer_churn.target_column,\n",
    "    class_labels=customer_churn.class_labels,\n",
    ")\n",
    "\n",
    "vm_train_ds = vm.init_dataset(\n",
    "    dataset=train_df,\n",
    "    input_id=\"train_dataset\",\n",
    "    target_column=customer_churn.target_column,\n",
    ")\n",
    "vm_train_ds.assign_predictions(model=vm_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run all the tests in the template, you can use the `vm.run_documentation_tests()` and pass the inputs we initialized above and the demo config from our customer_churn module. We will have to add a section to the config for our new test to tell it which inputs it should receive. This is done by simply adding a new element in the config dictionary where the key is the ID of the test and the value is a dictionary with the following structure:\n",
    "```python\n",
    "{\n",
    "    \"inputs\": {\n",
    "        \"model\": \"test_dataset\",\n",
    "        \"dataset\": \"model\",\n",
    "    }\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from validmind.utils import preview_test_config\n",
    "\n",
    "test_config = customer_churn.get_demo_test_config()\n",
    "test_config[\"my_custom_tests.ConfusionMatrix\"] = {\n",
    "    \"inputs\": {\n",
    "        \"dataset\": \"test_dataset\",\n",
    "        \"model\": \"model\",\n",
    "    }\n",
    "}\n",
    "preview_test_config(test_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_suite = vm.run_documentation_tests(config=test_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc9_'></a>\n",
    "\n",
    "## Some More Custom Tests\n",
    "\n",
    "Now that you understand the entire process of creating custom tests and using them in your documentation, let's create a few more to see different ways you can utilize custom tests."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc9_1_'></a>\n",
    "\n",
    "### Custom Test: Table of Model Hyperparameters\n",
    "\n",
    "This custom test will display a table of the hyperparameters used in the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@vm.test(\"my_custom_tests.Hyperparameters\")\n",
    "def hyperparameters(model):\n",
    "    \"\"\"The hyperparameters of a machine learning model are the settings that control the learning process.\n",
    "    These settings are specified before the learning process begins and can have a significant impact on the\n",
    "    performance of the model.\n",
    "\n",
    "    The hyperparameters of a model can be used to tune the model to achieve the best possible performance\n",
    "    on a given dataset. By examining the hyperparameters of a model, you can gain insight into how the model\n",
    "    was trained and how it might be improved.\n",
    "    \"\"\"\n",
    "    hyperparameters = model.model.get_xgb_params()  # dictionary of hyperparameters\n",
    "\n",
    "    # turn the dictionary into a table where each row contains a hyperparameter and its value\n",
    "    return [{\"Hyperparam\": k, \"Value\": v} for k, v in hyperparameters.items() if v]\n",
    "\n",
    "\n",
    "result = run_test(\"my_custom_tests.Hyperparameters\", inputs={\"model\": \"model\"})\n",
    "result.log()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the test has been run and logged, you can add it to your documentation using the same process as above. It should look like this:\n",
    "\n",
    "![screenshot showing hyperparameters test](../../images/hyperparameters-custom-metric.png)\n",
    "\n",
    "For our simple toy model, there are aren't really any proper hyperparameters but you can see how this could be useful for more complex models that have gone through hyperparameter tuning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc9_2_'></a>\n",
    "\n",
    "### Custom Test: External API Call\n",
    "\n",
    "This custom test will make an external API call to get the current BTC price and display it as a table. This demonstrates how you might integrate external data sources into your model documentation in a programmatic way. You could, for instance, setup a pipeline that runs a test like this every day to keep your model documentation in sync with an external system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import random\n",
    "\n",
    "\n",
    "@vm.test(\"my_custom_tests.ExternalAPI\")\n",
    "def external_api():\n",
    "    \"\"\"This test calls an external API to get a list of fake users. It then creates\n",
    "    a table with the relevant data so it can be displayed in the documentation.\n",
    "\n",
    "    The purpose of this test is to demonstrate how to call an external API and use the\n",
    "    data in a test. A test like this could even be setup to run in a scheduled\n",
    "    pipeline to keep your documentation in-sync with an external data source.\n",
    "    \"\"\"\n",
    "    url = \"https://jsonplaceholder.typicode.com/users\"\n",
    "    response = requests.get(url)\n",
    "    data = response.json()\n",
    "\n",
    "    # extract the time and the current BTC price in USD\n",
    "    return {\n",
    "        \"Model Owners/Stakeholders\": [\n",
    "            {\n",
    "                \"Name\": user[\"name\"],\n",
    "                \"Role\": random.choice([\"Owner\", \"Stakeholder\"]),\n",
    "                \"Email\": user[\"email\"],\n",
    "                \"Phone\": user[\"phone\"],\n",
    "                \"Slack Handle\": f\"@{user['name'].lower().replace(' ', '.')}\",\n",
    "            }\n",
    "            for user in data[:3]\n",
    "        ],\n",
    "        \"Model Developers\": [\n",
    "            {\n",
    "                \"Name\": user[\"name\"],\n",
    "                \"Role\": \"Developer\",\n",
    "                \"Email\": user[\"email\"],\n",
    "            }\n",
    "            for user in data[3:7]\n",
    "        ],\n",
    "        \"Model Validators\": [\n",
    "            {\n",
    "                \"Name\": user[\"name\"],\n",
    "                \"Role\": \"Validator\",\n",
    "                \"Email\": user[\"email\"],\n",
    "            }\n",
    "            for user in data[7:]\n",
    "        ],\n",
    "    }\n",
    "\n",
    "\n",
    "result = run_test(\"my_custom_tests.ExternalAPI\")\n",
    "result.log()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, you can add this to your documentation to see how it looks:\n",
    "\n",
    "![screenshot showing BTC price metric](../../images/external-data-custom-test.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc9_3_'></a>\n",
    "\n",
    "### Custom Test: Passing Parameters\n",
    "\n",
    "Custom test functions, as stated earlier, can take both inputs and params. When you define your function there is no need to distinguish between the two, the ValidMind Library will handle that for you. You simply need to add both to the function as arguments and the library will pass in the correct values.\n",
    "\n",
    "So for instance, if you wanted to parameterize the first custom test we created, the confusion matrix, you could do so like this:\n",
    "\n",
    "```python\n",
    "def confusion_matrix(dataset: VMDataset, model: VMModel, my_param: str = \"Default Value\"):\n",
    "    pass\n",
    "```\n",
    "\n",
    "And then when you run the test, you can pass in the parameter like this:\n",
    "\n",
    "```python\n",
    "vm.run_test(\n",
    "    \"my_custom_tests.ConfusionMatrix\",\n",
    "    inputs={\"model\": \"model\", \"dataset\": \"test_dataset\"},\n",
    "    params={\"my_param\": \"My Value\"},\n",
    ")\n",
    "```\n",
    "\n",
    "Or if you are running the entire documentation template, you would update the config like this:\n",
    "\n",
    "```python\n",
    "test_config[\"my_custom_tests.ConfusionMatrix\"] = {\n",
    "    \"inputs\": {\n",
    "        \"dataset\": \"test_dataset\",\n",
    "        \"model\": \"model\",\n",
    "    },\n",
    "    \"params\": {\n",
    "        \"my_param\": \"My Value\",\n",
    "    },\n",
    "}\n",
    "```\n",
    "\n",
    "Let's go ahead and create a toy test that takes a parameter and uses it in the result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly_express as px\n",
    "\n",
    "\n",
    "@vm.test(\"my_custom_tests.ParameterExample\")\n",
    "def parameter_example(\n",
    "    plot_title=\"Default Plot Title\", x_col=\"sepal_width\", y_col=\"sepal_length\"\n",
    "):\n",
    "    \"\"\"This test takes two parameters and creates a scatter plot based on them.\n",
    "\n",
    "    The purpose of this test is to demonstrate how to create a test that takes\n",
    "    parameters and uses them to generate a plot. This can be useful for creating\n",
    "    tests that are more flexible and can be used in a variety of scenarios.\n",
    "    \"\"\"\n",
    "    # return px.scatter(px.data.iris(), x=x_col, y=y_col, color=\"species\")\n",
    "    return px.scatter(\n",
    "        px.data.iris(), x=x_col, y=y_col, color=\"species\", title=plot_title\n",
    "    )\n",
    "\n",
    "\n",
    "result = run_test(\n",
    "    \"my_custom_tests.ParameterExample\",\n",
    "    params={\n",
    "        \"plot_title\": \"My Cool Plot\",\n",
    "        \"x_col\": \"sepal_width\",\n",
    "        \"y_col\": \"sepal_length\",\n",
    "    },\n",
    ")\n",
    "result.log()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Play around with this and see how you can use parameters, default values and other features to make your custom tests more flexible and useful.\n",
    "\n",
    "Here's how this one looks in the documentation:\n",
    "![screenshot showing parameterized test](../../images/parameterized-custom-metric.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc9_4_'></a>\n",
    "\n",
    "### Custom Test: Multiple Tables and Plots in a Single Test\n",
    "\n",
    "Custom test functions, as stated earlier, can return more than just one table or plot. In fact, any number of tables and plots can be returned. Let's see an example of this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import plotly_express as px\n",
    "\n",
    "\n",
    "@vm.test(\"my_custom_tests.ComplexOutput\")\n",
    "def complex_output():\n",
    "    \"\"\"This test demonstrates how to return many tables and figures in a single test\"\"\"\n",
    "    # create a couple tables\n",
    "    table = [{\"A\": 1, \"B\": 2}, {\"A\": 3, \"B\": 4}]\n",
    "    table2 = [{\"C\": 5, \"D\": 6}, {\"C\": 7, \"D\": 8}]\n",
    "\n",
    "    # create a few figures showing some random data\n",
    "    fig1 = px.line(x=np.arange(10), y=np.random.rand(10), title=\"Random Line Plot\")\n",
    "    fig2 = px.bar(x=[\"A\", \"B\", \"C\"], y=np.random.rand(3), title=\"Random Bar Plot\")\n",
    "    fig3 = px.scatter(\n",
    "        x=np.random.rand(10), y=np.random.rand(10), title=\"Random Scatter Plot\"\n",
    "    )\n",
    "\n",
    "    return (\n",
    "        {\n",
    "            \"My Cool Table\": table,\n",
    "            \"Another Table\": table2,\n",
    "        },\n",
    "        fig1,\n",
    "        fig2,\n",
    "        fig3,\n",
    "    )\n",
    "\n",
    "\n",
    "result = run_test(\"my_custom_tests.ComplexOutput\")\n",
    "result.log()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how you can return the tables as a dictionary where the key is the title of the table and the value is the table itself. You could also just return the tables by themselves but this way you can give them a title to more easily identify them in the result.\n",
    "\n",
    "![screenshot showing multiple tables and plots](../../images/multiple-tables-plots-custom-metric.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc9_5_'></a>\n",
    "\n",
    "### Custom Test: Images\n",
    "\n",
    "If you are using a plotting library that isn't supported by ValidMind (i.e. not `matplotlib` or `plotly`), you can still return the image directly as a bytes-like object. This could also be used to bring any type of image into your documentation in a programmatic way. For instance, you may want to include a diagram of your model architecture or a screenshot of a dashboard that your model is integrated with. As long as you can produce the image with Python or open it from a file, you can include it in your documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "@vm.test(\"my_custom_tests.Image\")\n",
    "def image():\n",
    "    \"\"\"This test demonstrates how to return an image in a test\"\"\"\n",
    "\n",
    "    # create a simple plot\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.plot([1, 2, 3, 4])\n",
    "    ax.set_title(\"Simple Line Plot\")\n",
    "\n",
    "    # save the plot as a PNG image (in-memory buffer)\n",
    "    img_data = io.BytesIO()\n",
    "    fig.savefig(img_data, format=\"png\")\n",
    "    img_data.seek(0)\n",
    "\n",
    "    plt.close()  # close the plot to avoid displaying it\n",
    "\n",
    "    return img_data.read()\n",
    "\n",
    "\n",
    "result = run_test(\"my_custom_tests.Image\")\n",
    "result.log()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adding this custom test to your documentation will display the image:\n",
    "\n",
    "![screenshot showing image custom test](../../images/image-in-custom-metric.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to log an image as a test result, you can do so by passing the path to the image as a parameter to the custom test and then opening the file in the test function. Here's an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@vm.test(\"my_custom_tests.MyPNGCorrelationMatrix\")\n",
    "def Image(path: str):\n",
    "    \"\"\"Opens a png image file and logs it as a test result to ValidMind\"\"\"\n",
    "    if not path.endswith(\".png\"):\n",
    "        raise ValueError(\"Image must be a PNG file\")\n",
    "\n",
    "    # return raw image bytes\n",
    "    with open(path, \"rb\") as f:\n",
    "        return f.read()\n",
    "    \n",
    "run_test(\n",
    "    \"my_custom_tests.MyPNGCorrelationMatrix\",\n",
    "    params={\"path\": \"../../images/pearson-correlation-matrix.png\"},\n",
    ").log()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The image is displayed in the test result:\n",
    "\n",
    "![screenshot showing image from file](../../images/pearson-correlation-matrix-test-output.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc9_6_'></a>\n",
    "\n",
    "### Custom Test: Description\n",
    "\n",
    "If you want to write a custom test description for your custom test instead of it is interpreted through llm, you can do so by returning string in your test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "@vm.test(\"my_custom_tests.MyCustomTest\")\n",
    "def my_custom_test(dataset, model):\n",
    "    \"\"\"\n",
    "    This is a custom computed test that computes confusion matrix for a binary classification model and return  a string as a test description.\n",
    "    \"\"\"\n",
    "    y_true = dataset.y\n",
    "    y_pred = dataset.y_pred(model)\n",
    "\n",
    "    confusion_matrix = metrics.confusion_matrix(y_true, y_pred)\n",
    "\n",
    "    cm_display = metrics.ConfusionMatrixDisplay(\n",
    "        confusion_matrix=confusion_matrix, display_labels=[False, True]\n",
    "    )\n",
    "    cm_display.plot()\n",
    "\n",
    "    plt.close()  # close the plot to avoid displaying it\n",
    "\n",
    "    return cm_display.figure_, \"Test Description - Confusion Matrix\", pd.DataFrame({\"Value\": [1, 2, 3]})  # return the figure object itself\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see here test result description has been customized here. The same result description will be displayed in the UI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = run_test(\n",
    "    \"my_custom_tests.MyCustomTest\",\n",
    "    inputs={\"model\": \"model\", \"dataset\": \"test_dataset\"},\n",
    ")\n",
    "result.log()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc10_'></a>\n",
    "\n",
    "## Conclusion\n",
    "\n",
    "In this notebook, we have demonstrated how to create custom tests in ValidMind. We have shown how to define custom test functions, register them with the ValidMind Library, run them against models and datasets, and add them to model documentation templates. We have also shown how to return tables and plots from custom tests and how to use them in the ValidMind Platform. We hope this tutorial has been helpful in understanding how to create and use custom tests in ValidMind."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc11_'></a>\n",
    "\n",
    "## Next steps\n",
    "\n",
    "You can look at the results of this test suite right in the notebook where you ran the code, as you would expect. But there is a better way — use the ValidMind Platform to work with your model documentation.\n",
    "\n",
    "<a id='toc11_1_'></a>\n",
    "\n",
    "### Work with your model documentation\n",
    "\n",
    "1. From the **Model Inventory** in the ValidMind Platform, go to the model you registered earlier. ([Need more help?](https://docs.validmind.ai/guide/model-inventory/working-with-model-inventory.html))\n",
    "\n",
    "2. Click and expand the **Model Development** section.\n",
    "\n",
    "What you see is the full draft of your model documentation in a more easily consumable version. From here, you can make qualitative edits to model documentation, view guidelines, collaborate with validators, and submit your model documentation for approval when it's ready. [Learn more ...](https://docs.validmind.ai/guide/model-documentation/working-with-model-documentation.html)\n",
    "\n",
    "<a id='toc11_2_'></a>\n",
    "\n",
    "### Discover more learning resources\n",
    "\n",
    "We offer many interactive notebooks to help you document models:\n",
    "\n",
    "- [Run tests & test suites](https://docs.validmind.ai/developer/model-testing/testing-overview.html)\n",
    "- [Code samples](https://docs.validmind.ai/developer/samples-jupyter-notebooks.html)\n",
    "\n",
    "Or, visit our [documentation](https://docs.validmind.ai/) to learn more about ValidMind."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc12_'></a>\n",
    "\n",
    "## Upgrade ValidMind\n",
    "\n",
    "<div class=\"alert alert-block alert-info\" style=\"background-color: #B5B5B510; color: black; border: 1px solid #083E44; border-left-width: 5px; box-shadow: 2px 2px 4px rgba(0, 0, 0, 0.2);border-radius: 5px;\">After installing ValidMind, you’ll want to periodically make sure you are on the latest version to access any new features and other enhancements.</div>\n",
    "\n",
    "Retrieve the information for the currently installed version of ValidMind:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip show validmind"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the version returned is lower than the version indicated in our [production open-source code](https://github.com/validmind/validmind-library/blob/prod/validmind/__version__.py), restart your notebook and run:\n",
    "\n",
    "```bash\n",
    "%pip install --upgrade validmind\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may need to restart your kernel after running the upgrade package for changes to be applied."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
