colabs/gemini_pro_vertex_evals_gsm8k_python.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "private_outputs": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/henrykmichalewski/math-evals/blob/master/colabs/gemini_pro_vertex_evals_gsm8k_python.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/henrykmichalewski/math-evals\n",
        "!pip install -e math-evals -q"
      ],
      "metadata": {
        "id": "8XfB110qFyL-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The above cell clones the unlocked variant of the OpenAI HumanEval repo.\n",
        "\n",
        "Unlocking means that the correctness code is actually executing completions. The repo contains some helper code to simplify gsm8k_python evaluations. You need to reset the runtime before moving forward."
      ],
      "metadata": {
        "id": "XUI6RP5buCGg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Loading gsm8k"
      ],
      "metadata": {
        "id": "jJH11qRe9Ax5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install datasets -q"
      ],
      "metadata": {
        "id": "PCe3l7wJu6ZI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset"
      ],
      "metadata": {
        "id": "9hfa5K9X8r0-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gsm8k = load_dataset(\"gsm8k\", \"main\")"
      ],
      "metadata": {
        "id": "DD07jtLW8xi6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gsm8k_train, gsm8k_test = gsm8k['train'], gsm8k['test']"
      ],
      "metadata": {
        "id": "QhpOyI1Q81K3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from human_eval.execution import check_gsm8k_correctness"
      ],
      "metadata": {
        "id": "CcGTboYNF2yt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check_gsm8k_correctness(gsm8k_test[5], \" return 63\", task_id=5)"
      ],
      "metadata": {
        "id": "AclQ1Z7jGP08"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check_gsm8k_correctness(gsm8k_test[5], \" return 64\", task_id=5)"
      ],
      "metadata": {
        "id": "x8221tb_GcKb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check_gsm8k_correctness(gsm8k_test[5], \" return 64.000001\", task_id=5)"
      ],
      "metadata": {
        "id": "EYsJAZ2DKXCl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Vertex API"
      ],
      "metadata": {
        "id": "d2AfwfJ0frRF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# I am running this from the colab terminal - for some reason the\n",
        "# !gcloud auth application-default login"
      ],
      "metadata": {
        "id": "hvBStAGVf6tf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip3 install --upgrade --user google-cloud-aiplatform"
      ],
      "metadata": {
        "id": "yVp5AICvf7dQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud auth application-default set-quota-project cloud-llm-preview2"
      ],
      "metadata": {
        "id": "CyZ4J6UTf8sW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import re\n",
        "import time\n",
        "from typing import Callable, Iterable, Match, Optional, Pattern, Protocol, Sequence, Union\n",
        "\n",
        "import vertexai"
      ],
      "metadata": {
        "id": "eMGkg-sDgGAK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from google.cloud.aiplatform_v1beta1.types import (\n",
        "    content as gapic_content_types,\n",
        ")\n",
        "## Add your GCP project and location data\n",
        "project_id = \"\"\n",
        "location = \"\"\n",
        "\n",
        "vertexai.init(project=project_id, location=location)\n",
        "\n",
        "from vertexai.preview.generative_models import GenerativeModel"
      ],
      "metadata": {
        "id": "y9TR0ZFhftQD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model_name = 'gemini-pro'"
      ],
      "metadata": {
        "id": "kqWUiZTbgbh5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = GenerativeModel(model_name)"
      ],
      "metadata": {
        "id": "eLnumr4lg5kG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Loading/Saving samples into gdrive"
      ],
      "metadata": {
        "id": "mHc56CE10RUH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "import os\n",
        "\n",
        "drive.mount('/content/gdrive/', force_remount=True)"
      ],
      "metadata": {
        "id": "482up10cQ_cL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "path = f'/content/gdrive/My Drive/vertex/{model_name}'\n",
        "os.makedirs(path, exist_ok=True)\n",
        "from human_eval.data import stream_jsonl, write_jsonl"
      ],
      "metadata": {
        "id": "tk2YisgJTSkC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Read jsonl"
      ],
      "metadata": {
        "id": "EGGFns126jYb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "try:\n",
        "  samples = list(stream_jsonl(os.path.join(path, \"samples_gsm8k.jsonl\")))\n",
        "except:\n",
        "  samples = []"
      ],
      "metadata": {
        "id": "8aBRab2o2qlR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(samples)"
      ],
      "metadata": {
        "id": "EbCVJ76Fb4AK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Helper function to generate vertex completions"
      ],
      "metadata": {
        "id": "om8kRwrmglsX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_one_completion(prompt):\n",
        "  while True:\n",
        "    try:\n",
        "      response = model.generate_content(\n",
        "        prompt,\n",
        "        generation_config={\n",
        "            'temperature' : 0.,\n",
        "            },\n",
        "      )\n",
        "      break\n",
        "    except Exception as e:\n",
        "      print(e)\n",
        "      print(\"Waiting for vertex...\")\n",
        "      time.sleep(30)\n",
        "  return response"
      ],
      "metadata": {
        "id": "vUnGb5SkgoEh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "generate_one_completion(\"Hi there, Gemini!\")"
      ],
      "metadata": {
        "id": "DxmjSBMIguvt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Full loop"
      ],
      "metadata": {
        "id": "sWbKV3tC7pN3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from data import gsm8k_python_prompt"
      ],
      "metadata": {
        "id": "rUsu9zF7vlPc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "all_correct = 0\n",
        "for task_id, problem in enumerate(gsm8k_test):\n",
        "\n",
        "    # Check if task_id is less than samples length\n",
        "    if task_id < len(samples):\n",
        "        correctness = check_gsm8k_correctness(problem, samples[task_id]['completion'], task_id=task_id)\n",
        "        correct = correctness['passed']\n",
        "        all_correct += correct\n",
        "        continue\n",
        "\n",
        "    # Print Task ID\n",
        "    print(f\"task_id {task_id}\")\n",
        "\n",
        "    # Formulate and print the full prompt\n",
        "    full_prompt = (gsm8k_python_prompt.PYTHON_PROMPT + '\\n' +\n",
        "                  gsm8k_python_prompt.PYTHON_TEMPLATE.format(question=problem['question']))\n",
        "\n",
        "    text_empty = False\n",
        "    while True:\n",
        "      try:\n",
        "        response = generate_one_completion(full_prompt)\n",
        "        try:\n",
        "          completion = response.text\n",
        "        except:\n",
        "          text_empty = True\n",
        "        break\n",
        "      except Exception as e:\n",
        "        print(e)\n",
        "        print(\"Waiting for vertex...\")\n",
        "        time.sleep(30)\n",
        "\n",
        "    if text_empty: continue\n",
        "\n",
        "    # Works for Pro, a model-dependent processing of outputs:\n",
        "    try:\n",
        "      completion = '\\n'.join(completion.split('```')[1].splitlines()[2:])\n",
        "    except:\n",
        "      print(\"Failed to extract program.\")\n",
        "\n",
        "    print(completion)\n",
        "\n",
        "    # New lines for visibility\n",
        "    print(\"\\n\")\n",
        "    print(problem['answer'])\n",
        "\n",
        "    # Check for correctness of the problem solved\n",
        "    correctness = check_gsm8k_correctness(problem, completion, task_id=task_id)\n",
        "    correct = correctness['passed']\n",
        "    all_correct += correct\n",
        "\n",
        "    # Print out the correctness\n",
        "    print(f\"Correctnes: {correct}\\nAll correct: {all_correct}\")\n",
        "    print(\"=\"*40)\n",
        "\n",
        "    # Append the task results to the list and check if need to write to JSON file\n",
        "    samples.append(dict(task_id=task_id, completion=completion))\n",
        "    if task_id % 100 == 0 and task_id:\n",
        "        write_jsonl(os.path.join(path, \"samples_gsm8k.jsonl\"), samples)\n",
        "\n",
        "# Write all samples to JSONL file at the end\n",
        "write_jsonl(os.path.join(path, \"samples_gsm8k.jsonl\"), samples)\n"
      ],
      "metadata": {
        "id": "4SnHfNaq5pyJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "1030 / 1319"
      ],
      "metadata": {
        "id": "JcEQ12Ea0Wg2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Write jsonl"
      ],
      "metadata": {
        "id": "ILUoZmKp6nAz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "write_jsonl(os.path.join(path, \"samples_gsm8k.jsonl\"), samples)"
      ],
      "metadata": {
        "id": "lA6qRPr6TQGC"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}