{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Visualization\n",
    "\n",
    "In this example, we adapt the [code](https://github.com/thunlp/MatPlotAgent) from the [MatPlotAgent paper](https://arxiv.org/abs/2402.11453).\n",
    "\n",
    "The workflow is as follows:\n",
    "- **Query expansion** adds more information or context regarding the user's task\n",
    "- **Initial code generation** generates plotting code based on the provided context\n",
    "- **Code execution** runs the code in a sandbox\n",
    "  - If the program has errors, it goes to a **debugging** agent, which will refine the code to remove errors and send it back for execution\n",
    "  - If the program does not have errors, it goes to a **vision feedback** agent that will provide feedback on the output plot. This feedback is then used to **refine** the code. \n",
    "\n",
    "Lastly, the refined code is executed and the final output image is generated.\n",
    "\n",
    "![datavis](../imgs/datavis.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Loader\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the benchmarking data provided by the paper, which you can download [here](https://github.com/thunlp/MatPlotAgent/tree/main/benchmark_data). For convenience, we've already split the dataset into train and test. You can find the exact splits in the `./benchmark_split/` folder.\n",
    "\n",
    "The dataset contains CSVs as data sources, instructions to generate a specific plot from a given data source, and ground-truth images. The workflow creates a temporary workspace for the sandbox. Hence, the input provides a `query` (instruction), `directory_path` (sandbox), `example_id` and `input_path` which are used to reference the training data and create the results folder. The `ground_truth` simply references the ground truth image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cognify\n",
    "import json\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "@cognify.register_data_loader\n",
    "def load_data():\n",
    "    def load_from_file(input_file):\n",
    "        # open the json file \n",
    "        data = json.load(open(input_file))\n",
    "        \n",
    "        all_data = []\n",
    "        for item in data:\n",
    "            novice_instruction = item['simple_instruction']\n",
    "            example_id = item['id']\n",
    "            directory_path = f'opt_runs'\n",
    "\n",
    "            if not os.path.exists(directory_path):\n",
    "                os.makedirs(directory_path, exist_ok=True)\n",
    "            \n",
    "            input = {\n",
    "                'query': novice_instruction,\n",
    "                \"directory_path\": directory_path,\n",
    "                \"example_id\": example_id,\n",
    "                \"input_path\": f'benchmark_data/data/{example_id}',\n",
    "            }\n",
    "            label = {\"ground_truth\": f\"benchmark_data/ground_truth/example_{example_id}.png\"}\n",
    "            all_data.append((input, label))\n",
    "        return all_data\n",
    "            \n",
    "    all_train = load_from_file('benchmark_split/train_data.json')\n",
    "    test_data = load_from_file('benchmark_split/test_data.json')\n",
    "    train_indices = np.random.choice(range(len(all_train)), 40, replace=False).tolist()\n",
    "    eval_indices = list(set(range(len(all_train))) - set(train_indices))\n",
    "    \n",
    "    train_data = [all_train[i] for i in train_indices]\n",
    "    eval_data = [all_train[i] for i in eval_indices]\n",
    "    return train_data, eval_data, test_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluator\n",
    "\n",
    "The evaluator calls GPT-4o to score the output with respect to the ground truth image. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "import warnings\n",
    "import re\n",
    "import base64\n",
    "\n",
    "BASE_URL='https://api.openai.com/v1'\n",
    "import dotenv\n",
    "dotenv.load_dotenv()\n",
    "API_KEY = os.environ[\"OPENAI_API_KEY\"]\n",
    "\n",
    "\n",
    "def encode_image(image_path):\n",
    "    with open(image_path, \"rb\") as image_file:\n",
    "        return base64.b64encode(image_file.read()).decode('utf-8')\n",
    "\n",
    "@cognify.register_evaluator\n",
    "def gpt_4o_evaluate(ground_truth, image, rollback):\n",
    "    client = OpenAI(\n",
    "        api_key=API_KEY,\n",
    "        base_url=BASE_URL,)\n",
    "    if not os.path.exists(f'{image}'):\n",
    "        if os.path.exists(f'{rollback}'):\n",
    "            base64_image1 = encode_image(f\"{ground_truth}\")\n",
    "            base64_image2 = encode_image(f\"{rollback}\")\n",
    "        else:\n",
    "            image = 'benchmark_data/ground_truth/empty.png'\n",
    "            base64_image1 = encode_image(f\"{image}\")\n",
    "            base64_image2 = encode_image(f\"{image}\")\n",
    "    else:\n",
    "        base64_image1 = encode_image(f\"{ground_truth}\")\n",
    "        base64_image2 = encode_image(f\"{image}\")\n",
    "\n",
    "    response = client.chat.completions.create(\n",
    "      model=\"gpt-4o\",\n",
    "      temperature=0.0,\n",
    "      messages=[\n",
    "        {\n",
    "          \"role\": \"user\",\n",
    "          \"content\": [\n",
    "            {\n",
    "              \"type\": \"text\",\n",
    "              \"text\": f'''You are an excellent judge at evaluating visualization plots between a model generated plot and the ground truth. You will be giving scores on how well it matches the ground truth plot.\n",
    "               \n",
    "               The generated plot will be given to you as the first figure. If the first figure is blank, that means the code failed to generate a figure.\n",
    "               Another plot will be given to you as the second figure, which is the desired outcome of the user query, meaning it is the ground truth for you to reference.\n",
    "               Please compare the two figures head to head and rate them.\n",
    "               Suppose the second figure has a score of 100, rate the first figure on a scale from 0 to 100.\n",
    "               Scoring should be carried out in the following aspect:\n",
    "               1. Plot correctness: \n",
    "               Compare closely between the generated plot and the ground truth, the more resemblance the generated plot has compared to the ground truth, the higher the score. The score should be proportionate to the resemblance between the two plots.\n",
    "               In some rare occurrence, see if the data points are generated randomly according to the query, if so, the generated plot may not perfectly match the ground truth, but it is correct nonetheless.\n",
    "               Only rate the first figure, the second figure is only for reference.\n",
    "               If the first figure is blank, that means the code failed to generate a figure. Give a score of 0 on the Plot correctness.\n",
    "                After scoring from the above aspect, please give a final score. The final score is preceded by the [FINAL SCORE] token.\n",
    "               For example [FINAL SCORE]: 40.''',\n",
    "            },\n",
    "            {\n",
    "              \"type\": \"image_url\",\n",
    "              \"image_url\": {\n",
    "                \"url\": f\"data:image/jpeg;base64,{base64_image2}\",\n",
    "              },\n",
    "            },\n",
    "            {\n",
    "              \"type\": \"image_url\",\n",
    "              \"image_url\": {\n",
    "                \"url\": f\"data:image/jpeg;base64,{base64_image1}\",\n",
    "              },\n",
    "            },\n",
    "          ],\n",
    "        }\n",
    "      ],\n",
    "      max_tokens=1000,\n",
    "    )\n",
    "    pattern = r'\\[FINAL SCORE\\]: (\\d{1,3})'\n",
    "\n",
    "\n",
    "    # Create a dictionary to store settings and their scores\n",
    "    match = re.search(pattern, response.choices[0].message.content)\n",
    "    if match:\n",
    "        score  = int(match.group(1)) / 100\n",
    "    else:\n",
    "        warnings.warn(\"No score found!!!\")\n",
    "        score = 0\n",
    "    return score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuring the Optimizer\n",
    "\n",
    "We've created a search option for data visualization that searches over the following:\n",
    "- Chain-of-Thought reasoning\n",
    "- Planning before acting\n",
    "- 2 few-shot examples\n",
    "- An ensemble of 3 agents for a task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ================= Configure Search =================\n",
    "from cognify.hub.search import datavis\n",
    "search_settings = datavis.create_search(opt_log_dir='opt_results', evaluator_batch_size=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Begin Optimization\n",
    "\n",
    "The above code blocks are also provided in `config.py`. You can use Cognify's CLI to start the optimization with\n",
    "```console\n",
    "$ cognify optimize workflow.py\n",
    "```\n",
    "\n",
    "Alternatively, you can run the following cell (*warning*: this workflow may run for quite some time):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, val, dev = load_data()\n",
    "\n",
    "opt_cost, pareto_frontier, opt_logs = cognify.optimize(\n",
    "    script_path=\"workflow.py\",\n",
    "    control_param=search_settings,\n",
    "    train_set=train,\n",
    "    val_set=val,\n",
    "    eval_fn=gpt_4o_evaluate,\n",
    "    force=True, # This will overwrite the existing results\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization Results\n",
    "\n",
    "Cognfiy will output each optimized workflow to a `.cog` file. For this workflow, the optimizer applies the following optimizations to specific agents:\n",
    "- add few-shot examples to the query expansion agent\n",
    "- ensemble the debugging agent\n",
    "  - add chain-of-thought to the first ensembled modules\n",
    "  - add chain-of-thought __and__ planning to the second ensembled module\n",
    "  - add few-shot examples to the third ensembled module\n",
    "- ensemble the vision-based feedback agent\n",
    "  - add few-shot examples to all the ensembled modules\n",
    "  - add chain-of-thought to the first ensembled module\n",
    "  - add chain-of-thought __and__ planning to the second ensembled module\n",
    "- use few-shot examples and planning for code refinement\n",
    "\n",
    "The final optimized workflow is depicted below, with optimizations highlighted in green.\n",
    "\n",
    "![datavis-opt](../imgs/datavis_optimized.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check out more details on [how to interpret optimization results](https://cognify-ai.readthedocs.io/en/latest/user_guide/tutorials/interpret.html#detailed-transformation-trace)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chess-ex",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}