{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Document Retrieval\n",
    "\n",
    "In this example, we build a three-hop document retrival workflow for [HoVer](https://arxiv.org/abs/2011.03088). The workflow gathers related context to verify a given claim. We aim to improve the f1-score over the retrieved documents.\n",
    "\n",
    "The workflow two types of agents (4 in total):\n",
    "- **Query agent**: generates a search query for retrieval.\n",
    "- **Summarize agent**: extracts useful points in the given documents.\n",
    "\n",
    "![hover](../imgs/hover.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) Setup\n",
    "\n",
    "First, let's set the environment for workflow execution. Following keys are required:\n",
    "\n",
    "- OPENAI_API_KEY=\"your-openai-key\"\n",
    "- COLBERT_URL=\"colbert-serving-url\"\n",
    "\n",
    "> **Note:** \n",
    ">\n",
    "> If you are using DSPy's ColBERT service, try link `http://20.102.90.50:2017/wiki17_abstracts`. \n",
    ">\n",
    "> For hosting on your local machine, check [ColBERT official repo](https://github.com/stanford-futuredata/ColBERT) for installation and setup."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2) Check Hover Workflow\n",
    "\n",
    "The complete code for this workflow is based on `dspy` and is avaibale in `workflow.py`. \n",
    "\n",
    "The workflow returns the `pid` of all retrieved passages, which will be used to calcuate the f1-score.\n",
    "\n",
    "Try it out with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'pred_docs': (670219, 821694, 1513707, 4434531, 4751601, 670219, 813146, 821694, 821694, 1722727)}\n"
     ]
    }
   ],
   "source": [
    "%run workflow.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3) Optimize The Workflow\n",
    "\n",
    "The workflow entry point is already registered using annotation `cognify.register_workflow`.\n",
    "\n",
    "Here we configure the optimization pipeline:\n",
    "1. Define the evaluation method\n",
    "2. Define the data loader\n",
    "3. Config the optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Use F1 as the retrieval metric\n",
    "\n",
    "We use builtin `f1_score_set` to measure the overlap between retrieved documents and the ground truth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cognify\n",
    "from cognify.hub.evaluators import f1_score_set\n",
    "\n",
    "@cognify.register_evaluator\n",
    "def doc_f1(pred_docs, gold_docs):\n",
    "    pred_docs = set(pred_docs)\n",
    "    gold_docs = set(gold_docs)\n",
    "    return f1_score_set(pred_docs, gold_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Load the Data\n",
    "\n",
    "We provide the example data in `qas._json` file for you to start with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"qid\":0,\"question\":\"Skagen Painter Peder Severin Kr\\u00f8yer favored naturalism along with Theodor Esbern Philipsen and the artist Ossian Elgstr\\u00f6m studied with in the early 1900s.\",\"support_pids\":[670219,1513707,811545],\"support_facts\":[[670219,0],[670219,1],[1513707,1],[811545,2]],\"support_titles\":[\"Kristian Zahrtmann\",\"Peder Severin Kr\\u00f8yer\",\"Ossian Elgstr\\u00f6m\"],\"num_hops\":3,\"label\":1,\"uid\":\"330ca632-e83f-4011-b11b-0d0158145036\",\"hpqa_id\":\"5ab7a86d5542995dae37e986\"}\n"
     ]
    }
   ],
   "source": [
    "!head -n 1 qas._json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data should be formatted to align with the function signature of both the workflow entry point and the evaluator.\n",
    "\n",
    "Signatures are:\n",
    "- workflow (claim) -> {'pred_docs': ...}\n",
    "- evaluator (pred_docs, gold_docs) -> float\n",
    "\n",
    "Thus the dataloader should provide a tuple of dictionaries like this:\n",
    "```python\n",
    "(\n",
    "    {'claim': ...},\n",
    "    {'gold_docs': ...},\n",
    ")\n",
    "```\n",
    "\n",
    "The complete data loader code is provided below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_size = 100\n",
    "val_size = 50\n",
    "dev_size = 200\n",
    "data_path = 'qas._json'\n",
    "seed = 0\n",
    "\n",
    "import json\n",
    "import random\n",
    "\n",
    "@cognify.register_data_loader\n",
    "def load_data():\n",
    "    data = []\n",
    "    \n",
    "    # only include difficult examples with 3 hops\n",
    "    with open(data_path, 'r') as file:\n",
    "        for line in file:\n",
    "            obj = json.loads(line)\n",
    "            if obj['num_hops'] == 3:\n",
    "                data.append(obj)\n",
    "                \n",
    "    rng = random.Random(seed)\n",
    "    rng.shuffle(data)\n",
    "    \n",
    "    def formatting(x):\n",
    "        input = {'claim': x['question']}\n",
    "        ground_truth = {'gold_docs': x['support_pids']}\n",
    "        return (input, ground_truth)\n",
    "    \n",
    "    train_set = data[:train_size]\n",
    "    val_set = data[train_size:train_size+val_size]\n",
    "    dev_set = data[train_size+val_size:train_size+val_size+dev_size]\n",
    "    \n",
    "    train_set = [formatting(x) for x in train_set]\n",
    "    val_set = [formatting(x) for x in val_set]\n",
    "    dev_set = [formatting(x) for x in dev_set]\n",
    "    return train_set, val_set, dev_set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Config the optimizer\n",
    "\n",
    "Let's use the default configuration to optimize this workflow. The search space includes:\n",
    "- 2 fewshot examples to add for each agent\n",
    "- whether to apply Chain-of-thought to each agent\n",
    "\n",
    "This workflow invokes many agents, we set the parallel level to `50` to accelerate the process. Please check your OpenAI rate limit and set accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cognify.hub.search import default\n",
    "\n",
    "search_settings = default.create_search(evaluator_batch_size=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4) Start the Optimization\n",
    "\n",
    "You can save the above configs in `config.py` file and use Cognify's CLI to fire the optimization with:\n",
    "\n",
    "```console\n",
    "$ cognify optimize workflow.py\n",
    "```\n",
    "\n",
    "Alternatively you can run the following:"
   ]
  },
  {
   "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=doc_f1,\n",
    "    force=True, # This will overwrite the existing results\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5) 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",
    "- ensemble both summarizers\n",
    "  - for the first summarizer, use chain-of-thought for two of the ensembled modules\n",
    "  - for the second summarizer, use chain-of-thought for one of the ensembled modules\n",
    "- ensemble the first query generation module\n",
    "  - use chain-of-thought for two of the ensembled modules\n",
    "- use chain-of-thought for the final query generation module\n",
    "\n",
    "The final optimized workflow is depicted below, with optimizations highlighted in green.\n",
    "\n",
    "![hover-opt](../imgs/hover_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)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fresh_env",
   "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
}