{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5339272e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-03-22T17:00:42.285169Z",
     "iopub.status.busy": "2023-03-22T17:00:42.284973Z",
     "iopub.status.idle": "2023-03-22T17:00:42.291455Z",
     "shell.execute_reply": "2023-03-22T17:00:42.290871Z"
    },
    "tags": [
     "hide_inp"
    ]
   },
   "outputs": [],
   "source": [
    "desc = \"\"\"\n",
    "### Book QA\n",
    "\n",
    "Chain that does question answering with Hugging Face embeddings. [[Code](https://github.com/srush/MiniChain/blob/main/examples/gatsby.py)]\n",
    "\n",
    "(Adapted from the [LlamaIndex example](https://github.com/jerryjliu/gpt_index/blob/main/examples/gatsby/TestGatsby.ipynb).)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f966010d",
   "metadata": {},
   "source": [
    "$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5391c476",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-03-22T17:00:42.293927Z",
     "iopub.status.busy": "2023-03-22T17:00:42.293738Z",
     "iopub.status.idle": "2023-03-22T17:00:43.695402Z",
     "shell.execute_reply": "2023-03-22T17:00:43.694722Z"
    }
   },
   "outputs": [],
   "source": [
    "import datasets\n",
    "import numpy as np\n",
    "from minichain import prompt, show, HuggingFaceEmbed, OpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b1b2a82",
   "metadata": {},
   "source": [
    "Load data with embeddings (computed beforehand)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a54cf84e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-03-22T17:00:43.698121Z",
     "iopub.status.busy": "2023-03-22T17:00:43.697792Z",
     "iopub.status.idle": "2023-03-22T17:00:43.730349Z",
     "shell.execute_reply": "2023-03-22T17:00:43.729747Z"
    },
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "86d0a2ceb7ad4f99978e37c2719f2960",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['passages', 'embeddings'],\n",
       "    num_rows: 52\n",
       "})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gatsby = datasets.load_from_disk(\"gatsby\")\n",
    "gatsby.add_faiss_index(\"embeddings\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9a08061",
   "metadata": {},
   "source": [
    "Fast KNN retieval prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3e69cd6b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-03-22T17:00:43.735462Z",
     "iopub.status.busy": "2023-03-22T17:00:43.735139Z",
     "iopub.status.idle": "2023-03-22T17:00:43.738420Z",
     "shell.execute_reply": "2023-03-22T17:00:43.737964Z"
    },
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "@prompt(HuggingFaceEmbed(\"sentence-transformers/all-mpnet-base-v2\"))\n",
    "def get_neighbors(model, inp, k=1):\n",
    "    embedding = model(inp)\n",
    "    res = olympics.get_nearest_examples(\"embeddings\", np.array(embedding), k)\n",
    "    return res.examples[\"passages\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "14e22d0a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-03-22T17:00:43.740824Z",
     "iopub.status.busy": "2023-03-22T17:00:43.740339Z",
     "iopub.status.idle": "2023-03-22T17:00:43.743342Z",
     "shell.execute_reply": "2023-03-22T17:00:43.742905Z"
    },
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "@prompt(OpenAI(),\n",
    "        template_file=\"gatsby.pmpt.tpl\")\n",
    "def ask(model, query, neighbors):\n",
    "    return model(dict(question=query, docs=neighbors))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bfca0bea",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-03-22T17:00:43.745377Z",
     "iopub.status.busy": "2023-03-22T17:00:43.745056Z",
     "iopub.status.idle": "2023-03-22T17:00:43.747768Z",
     "shell.execute_reply": "2023-03-22T17:00:43.747352Z"
    },
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def gatsby(query):\n",
    "    n = get_neighbors(query)\n",
    "    return ask(query, n)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "159b0b85",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8e3f74d0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-03-22T17:00:43.749935Z",
     "iopub.status.busy": "2023-03-22T17:00:43.749750Z",
     "iopub.status.idle": "2023-03-22T17:00:44.094814Z",
     "shell.execute_reply": "2023-03-22T17:00:44.094179Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7861\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7861/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gradio = show(gatsby,\n",
    "              subprompts=[get_neighbors, ask],\n",
    "              examples=[\"What did Gatsby do before he met Daisy?\",\n",
    "                        \"What did the narrator do after getting back to Chicago?\"],\n",
    "              keys={\"HF_KEY\"},\n",
    "              description=desc,\n",
    "              code=open(\"gatsby.py\", \"r\").read().split(\"$\")[1].strip().strip(\"#\").strip()\n",
    "              )\n",
    "if __name__ == \"__main__\":\n",
    "    gradio.launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdbe3ca1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7e96a1dc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-03-22T17:00:44.097451Z",
     "iopub.status.busy": "2023-03-22T17:00:44.097193Z",
     "iopub.status.idle": "2023-03-22T17:00:44.100135Z",
     "shell.execute_reply": "2023-03-22T17:00:44.099538Z"
    },
    "tags": [
     "hide_inp"
    ]
   },
   "outputs": [],
   "source": [
    "# QAPrompt().show({\"question\": \"Who was Gatsby?\", \"docs\": [\"doc1\", \"doc2\", \"doc3\"]}, \"\")\n",
    "# # -\n",
    "\n",
    "# show_log(\"gatsby.log\")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "tags,-all"
  },
  "kernelspec": {
   "display_name": "minichain",
   "language": "python",
   "name": "minichain"
  },
  "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.6"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "044c0c7d8f264a629d9a834ad7a3d02b": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_71c01d11477848a0a65eb6d8fa5c290c",
       "max": 1.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_b2191e94a5124ecaa10fcbba69552cad",
       "tabbable": null,
       "tooltip": null,
       "value": 1.0
      }
     },
     "0d3429e6725e49a0a63d578b0b46a1b8": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "2d0cb296ceba4f28afac033d606f3f80": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "306271182a3b4f51bfd69ea6467b7394": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "71c01d11477848a0a65eb6d8fa5c290c": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "86d0a2ceb7ad4f99978e37c2719f2960": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_f4226b39d8114ae291c8c1f4150b97b3",
        "IPY_MODEL_044c0c7d8f264a629d9a834ad7a3d02b",
        "IPY_MODEL_cf381c0f6e0d41408fdff78193c8e003"
       ],
       "layout": "IPY_MODEL_0d3429e6725e49a0a63d578b0b46a1b8",
       "tabbable": null,
       "tooltip": null
      }
     },
     "b2191e94a5124ecaa10fcbba69552cad": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "bb8de7546e2d4a3d8c4b53a3170a9aa3": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "cf381c0f6e0d41408fdff78193c8e003": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_2d0cb296ceba4f28afac033d606f3f80",
       "placeholder": "​",
       "style": "IPY_MODEL_e239bdcb8e0f48f9b0cec18f6280bd8a",
       "tabbable": null,
       "tooltip": null,
       "value": " 1/1 [00:00&lt;00:00, 101.96it/s]"
      }
     },
     "e239bdcb8e0f48f9b0cec18f6280bd8a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "f4226b39d8114ae291c8c1f4150b97b3": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_306271182a3b4f51bfd69ea6467b7394",
       "placeholder": "​",
       "style": "IPY_MODEL_bb8de7546e2d4a3d8c4b53a3170a9aa3",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}