py/cookbooks/bootstrap_few_shot/lunary.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Import the Zenbase Library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "import subprocess\n",
    "\n",
    "def install_package(package):\n",
    "    try:\n",
    "        subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", package])\n",
    "    except subprocess.CalledProcessError as e:\n",
    "        print(f\"Failed to install {package}: {e}\")\n",
    "        raise\n",
    "\n",
    "def install_packages(packages):\n",
    "    for package in packages:\n",
    "        install_package(package)\n",
    "\n",
    "try:\n",
    "    # Check if running in Google Colab\n",
    "    import google.colab\n",
    "    IN_COLAB = True\n",
    "except ImportError:\n",
    "    IN_COLAB = False\n",
    "\n",
    "if IN_COLAB:\n",
    "    # Install the zenbase package if running in Google Colab\n",
    "    # install_package('zenbase')\n",
    "    # Install the zenbse package from a GitHub branch if running in Google Colab\n",
    "    install_package('git+https://github.com/zenbase-ai/lib.git@main#egg=zenbase&subdirectory=py')\n",
    "\n",
    "    # List of other packages to install in Google Colab\n",
    "    additional_packages = [\n",
    "        'python-dotenv',\n",
    "        'lunary',\n",
    "        'openai',\n",
    "        'langchain',\n",
    "        'langchain_openai'\n",
    "    ]\n",
    "    \n",
    "    # Install additional packages\n",
    "    install_packages(additional_packages)\n",
    "\n",
    "# Now import the zenbase library\n",
    "try:\n",
    "    import zenbase\n",
    "except ImportError as e:\n",
    "    print(\"Failed to import zenbase: \", e)\n",
    "    raise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Configure the Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# import os\n",
    "#\n",
    "# os.environ[\"OPENAI_API_KEY\"] = \"...\"\n",
    "# os.environ[\"LUNARY_PUBLIC_KEY\"] = \"...\"\n",
    "\n",
    "load_dotenv(Path(\"../../.env.test\"), override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Initial Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "import lunary\n",
    "\n",
    "openai = OpenAI()\n",
    "lunary.monitor(openai)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Now, you probably already have some LLM code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "It could use the OpenAI SDK, LangChain, or anything really. But it looks something like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import openai\n",
    "from typing import Dict\n",
    "\n",
    "def solver(inputs: str) -> str:\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": \"You are an expert math solver. Provide the numerical answer to the given question based on the provided plan and operation. Return only the number as a JSON object in the format: {\\\"answer\\\": \\\"<number>\\\"}\"\n",
    "        }\n",
    "    ]\n",
    "\n",
    "    plan = planner_chain(inputs)\n",
    "    operation = operation_finder({\"plan\": plan[\"plan\"], \"question\": inputs})\n",
    "\n",
    "    messages.extend([\n",
    "        {\"role\": \"user\", \"content\": f\"Question: {inputs}\"},\n",
    "        {\"role\": \"user\", \"content\": f\"Plan: {plan['plan']}\"},\n",
    "        {\"role\": \"user\", \"content\": f\"Mathematical Operation needed: {operation['operation']}\"},\n",
    "        {\"role\": \"user\", \"content\": \"Provide only the numerical answer.\"}\n",
    "    ])\n",
    "\n",
    "    response = openai.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=messages,\n",
    "        response_format={\"type\": \"json_object\"},\n",
    "    )\n",
    "\n",
    "    answer = json.loads(response.choices[0].message.content)\n",
    "    return answer[\"answer\"]\n",
    "\n",
    "def planner_chain(inputs: str) -> Dict[str, str]:\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": \"You are an expert math solver. Create a step-by-step plan to solve the given question. Return the plan as a JSON object in the format: {\\\"plan\\\": \\\"<step-by-step plan>\\\"}\"\n",
    "        },\n",
    "        {\"role\": \"user\", \"content\": inputs}\n",
    "    ]\n",
    "\n",
    "    response = openai.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=messages,\n",
    "        response_format={\"type\": \"json_object\"},\n",
    "    )\n",
    "\n",
    "    return json.loads(response.choices[0].message.content)\n",
    "\n",
    "def operation_finder(inputs: Dict[str, str]) -> Dict[str, str]:\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": \"You are an expert math solver. Identify the primary mathematical operation needed to solve the problem based on the given plan. Use simple operations like addition, subtraction, multiplication, or division. Return the operation as a JSON object in the format: {\\\"operation\\\": \\\"<operation>\\\"}\"\n",
    "        },\n",
    "        {\"role\": \"user\", \"content\": f\"Question: {inputs['question']}\"},\n",
    "        {\"role\": \"user\", \"content\": f\"Plan: {inputs['plan']}\"}\n",
    "    ]\n",
    "\n",
    "    response = openai.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=messages,\n",
    "        response_format={\"type\": \"json_object\"},\n",
    "    )\n",
    "\n",
    "    return json.loads(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Test the function to see if that works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "solver(\"What is 2 + 2?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Then you're probably evaluating like this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import lunary \n",
    "evalset = lunary.get_dataset(\"gsm8k-evalset\")\n",
    "\n",
    "scores = []\n",
    "for item in evalset:\n",
    "    answer = solver(item.input)\n",
    "    passed, results = lunary.evaluate(\n",
    "        \"exact-match\",\n",
    "        input=item.input,\n",
    "        output=answer,\n",
    "        ideal_output=item.ideal_output.split(\"#### \")[-1],\n",
    "    )\n",
    "    scores.append(int(passed))\n",
    "\n",
    "print(\"Average score\", sum(scores) / len(scores))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    " # Now, how can we optimize this score?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## First, initialize the Zenbase ZenbaseTracer and ZenLunary objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from zenbase.core.managers import ZenbaseTracer\n",
    "from zenbase.adaptors.lunary import ZenLunary\n",
    "\n",
    "zenbase_tracer = ZenbaseTracer()\n",
    "lunary_adaptor = ZenLunary(client=lunary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from typing import Dict, Any\n",
    "from zenbase.types import LMRequest\n",
    "import openai\n",
    "\n",
    "@zenbase_tracer\n",
    "def solver(request: LMRequest) -> str:\n",
    "    system_message = {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": (\n",
    "            \"You are an expert math solver. Analyze the given question, follow the provided step-by-step plan, \"\n",
    "            \"and perform the specified mathematical operation. Your response should be a single number representing \"\n",
    "            \"the final answer. Format your response as a JSON object: {\\\"answer\\\": \\\"[numerical result]\\\"}\"\n",
    "        )\n",
    "    }\n",
    "    \n",
    "    messages = [system_message]\n",
    "    \n",
    "    for demo in request.zenbase.task_demos:\n",
    "        messages.extend([\n",
    "            {\"role\": \"user\", \"content\": f\"Example Question: {demo.inputs}\"},\n",
    "            {\"role\": \"assistant\", \"content\": f\"Example Answer: {demo.outputs}\"}\n",
    "        ])\n",
    "\n",
    "    plan = planner_chain(request.inputs)\n",
    "    operation = operation_finder({\"plan\": plan[\"plan\"], \"question\": request.inputs})\n",
    "\n",
    "    messages.extend([\n",
    "        {\"role\": \"user\", \"content\": f\"Question: {request.inputs}\"},\n",
    "        {\"role\": \"user\", \"content\": f\"Plan: {plan['plan']}\"},\n",
    "        {\"role\": \"user\", \"content\": f\"Required Operation: {operation['operation']}\"},\n",
    "        {\"role\": \"user\", \"content\": \"Provide only the numerical answer in the specified JSON format.\"}\n",
    "    ])\n",
    "\n",
    "    response = openai.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=messages,\n",
    "        response_format={\"type\": \"json_object\"},\n",
    "    )\n",
    "\n",
    "    return json.loads(response.choices[0].message.content)[\"answer\"]\n",
    "\n",
    "@zenbase_tracer\n",
    "def planner_chain(request: LMRequest) -> Dict[str, str]:\n",
    "    system_message = {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": (\n",
    "            \"As an expert math solver, create a detailed step-by-step plan to solve the given problem. \"\n",
    "            \"Your plan should be clear, concise, and easy to follow. \"\n",
    "            \"Format your response as a JSON object: {\\\"plan\\\": \\\"[step-by-step plan]\\\"}\"\n",
    "        )\n",
    "    }\n",
    "    \n",
    "    messages = [system_message]\n",
    "    \n",
    "    if request.zenbase.task_demos:\n",
    "        for demo in request.zenbase.task_demos[:2]:\n",
    "            messages.extend([\n",
    "                {\"role\": \"user\", \"content\": demo.inputs},\n",
    "                {\"role\": \"assistant\", \"content\": demo.outputs[\"plan\"]}\n",
    "            ])\n",
    "    \n",
    "    messages.append({\"role\": \"user\", \"content\": request.inputs})\n",
    "\n",
    "    response = openai.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=messages,\n",
    "        response_format={\"type\": \"json_object\"},\n",
    "    )\n",
    "\n",
    "    return json.loads(response.choices[0].message.content)\n",
    "\n",
    "@zenbase_tracer\n",
    "def operation_finder(request: LMRequest) -> Dict[str, str]:\n",
    "    system_message = {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": (\n",
    "            \"Based on the given question and solution plan, identify the primary mathematical operation required. \"\n",
    "            \"Focus on basic operations: addition, subtraction, multiplication, or division. \"\n",
    "            \"Format your response as a JSON object: {\\\"operation\\\": \\\"[primary operation]\\\"}\"\n",
    "        )\n",
    "    }\n",
    "    \n",
    "    messages = [system_message]\n",
    "    \n",
    "    if request.zenbase.task_demos:\n",
    "        for demo in request.zenbase.task_demos[:2]:\n",
    "            messages.extend([\n",
    "                {\"role\": \"user\", \"content\": f\"Question: {demo.inputs['question']}\"},\n",
    "                {\"role\": \"user\", \"content\": f\"Plan: {demo.inputs['plan']}\"},\n",
    "                {\"role\": \"assistant\", \"content\": demo.outputs[\"operation\"]}\n",
    "            ])\n",
    "    \n",
    "    messages.extend([\n",
    "        {\"role\": \"user\", \"content\": f\"Question: {request.inputs['question']}\"},\n",
    "        {\"role\": \"user\", \"content\": f\"Plan: {request.inputs['plan']}\"}\n",
    "    ])\n",
    "\n",
    "    response = openai.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=messages,\n",
    "        response_format={\"type\": \"json_object\"},\n",
    "    )\n",
    "\n",
    "    return json.loads(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Now we can optimize!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Set up your optimizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "evaluator_kwargs = dict(\n",
    "    checklist=\"exact-match\",\n",
    "    concurrency=2,\n",
    ")\n",
    "\n",
    "# for lunary there is not feature to create dataset with code, so dataset are created\n",
    "# manually with UI, if you want to replicate the test on your own, you should put\n",
    "# GSM8K examples to dataset name like below:\n",
    "TRAIN_SET = \"gsmk8k-train-set\"\n",
    "TEST_SET = \"gsm8k-test-set\"\n",
    "VALIDATION_SET = \"gsm8k-validation-set\"\n",
    "\n",
    "assert lunary_adaptor.fetch_dataset_demos(TRAIN_SET) is not None\n",
    "assert lunary_adaptor.fetch_dataset_demos(TEST_SET) is not None\n",
    "assert lunary_adaptor.fetch_dataset_demos(VALIDATION_SET) is not None\n",
    "\n",
    "SAMPLES = 2\n",
    "SHOTS = 3\n",
    "\n",
    "from zenbase.optim.metric.bootstrap_few_shot import BootstrapFewShot\n",
    "\n",
    "bootstrap_few_shot = BootstrapFewShot(\n",
    "    shots=SHOTS,\n",
    "    training_set=TRAIN_SET,\n",
    "    test_set=TEST_SET,\n",
    "    validation_set=VALIDATION_SET,\n",
    "    evaluator_kwargs=evaluator_kwargs,\n",
    "    zen_adaptor=lunary_adaptor,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Do the optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_fn, _ = bootstrap_few_shot.perform(\n",
    "    solver,\n",
    "    samples=SAMPLES,\n",
    "    rounds=1,\n",
    "    trace_manager=zenbase_tracer,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Introspect evaluation improvement\n",
    "\n",
    "You can see in this example that the best function has improved the evaluation score by 50%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "bootstrap_few_shot.base_evaluation.evals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "bootstrap_few_shot.best_evaluation.evals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Use your optimized function\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "zenbase_tracer.flush()"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_fn(\"If I have 30% of shares, and Mo has 24.5% of shares, how many of our 10M shares are unassigned?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Introspect function traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function_traces = [v for k, v in zenbase_tracer.all_traces.items()][0][\"optimized\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Check the optimized parameters for solver\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "pprint(function_traces[\"solver\"][\"args\"][\"request\"].zenbase.task_demos)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Check the optimized parameters for planner_chain\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "pprint(function_traces[\"planner_chain\"][\"args\"][\"request\"].zenbase.task_demos)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Check the optimized parameters for operation_finder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "pprint(function_traces[\"operation_finder\"][\"args\"][\"request\"].zenbase.task_demos)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## How to save the function and load it later\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Save the optimized function args to a file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "bootstrap_few_shot.save_optimizer_args(\"bootstrap_few_shot_args.zenbase\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Load the optimized function args with the function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "optimized_function = bootstrap_few_shot.load_optimizer_and_function(\"bootstrap_few_shot_args.zenbase\", solver, zenbase_tracer)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Use the loaded function and make sure it loaded the demos.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "zenbase_tracer.flush()\n",
    "optimized_function(\"If I have 30% of shares, and Mo has 24.5% of shares, how many of our 10M shares are unassigned?\")\n",
    "function_traces = [v for k, v in zenbase_tracer.all_traces.items()][0][\"optimized\"]\n",
    "from pprint import pprint\n",
    "\n",
    "pprint(function_traces[\"solver\"][\"args\"][\"request\"].zenbase.task_demos)\n",
    "pprint(function_traces[\"planner_chain\"][\"args\"][\"request\"].zenbase.task_demos)\n",
    "pprint(function_traces[\"operation_finder\"][\"args\"][\"request\"].zenbase.task_demos)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}