examples/10K_Evaluator/10K_PDF_Evaluator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7cbc4c4a",
   "metadata": {},
   "source": [
    "# Build Your Own 10K Agent: Transform a Unstructured Financial Report to an Finetuned LLM\n",
    "Do you want to build an agent so that you can ask it anything about the annual report (10K)? In this example, we will show you how use `uniflow` and `pykoi` to extract knowledge from a unstructured annual report (10K) and then finetune an LLM on these knowledge.\n",
    "\n",
    "First, we'll use `uniflow` to generate question-answers (QAs) from a pdf using OpenAI's models via `uniflow`'s `MultiFlowPipeline`.\n",
    "\n",
    "Next, we'll use `pykoi` to run supervised fine-tuning (SFT) on the QAs generated by `uniflow`.\n",
    "\n",
    "Finally, we'll use `pykoi`'s Chatbot to run the SFT model, so you can ask questions about the 10K and get answers.\n",
    "\n",
    "For this example, we're using a 10K from [Nike](https://investors.nike.com/investors/news-events-and-reports/), [Amazon](https://ir.aboutamazon.com/sec-filings/sec-filings-details/default.aspx?FilingId=16361618), and [Alphabet](https://abc.xyz/investor/sec-filings/annual-filings/2023/).\n",
    "\n",
    ">*Note: In order to run this notebook, you need a GPU (for the `RLHF`).*\n",
    "\n",
    "### Before running the code\n",
    "\n",
    "You will need to set up a conda environment to run this notebook. You can set up the environment following the [instruction](https://github.com/CambioML/cambio-recipes/tree/main#installation).\n",
    "\n",
    "We are using uniflow and several of the pykoi modules, so you will need to install these in your environment as well:\n",
    "```\n",
    "pip3 install uniflow\n",
    "pip3 install \"pykoi[huggingface, rag, rlhf]\"\n",
    "```\n",
    "Finally, you will need to install torch:\n",
    "```\n",
    "pip3 uninstall torch\n",
    "pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121  # cu121 means cuda 12.1\n",
    "```\n",
    "\n",
    "Next, you will need a valid [OpenAI API key](https://platform.openai.com/api-keys) to run the code. Once you have the key, set it as the environment variable `OPENAI_API_KEY` within a `.env` file in the root directory of this repository. For more details, see this [instruction](https://github.com/CambioML/cambio-recipes/tree/main#api-keys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Generate QAs from a 10K using `uniflow`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Update System Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "\n",
    "sys.path.append(\".\")\n",
    "sys.path.append(\"..\")\n",
    "sys.path.append(\"../..\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install helper packages\n",
    "If you already have these installed, feel free to skip this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: ...\n"
     ]
    }
   ],
   "source": [
    "!{sys.executable} -m pip install pandas nougat-ocr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Dependency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8d84dd70",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/cambio-recipes-10k/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "from uniflow.pipeline import MultiFlowsPipeline\n",
    "from uniflow.flow.config import PipelineConfig\n",
    "from uniflow.flow.config import TransformOpenAIConfig, ExtractPDFConfig\n",
    "from uniflow.flow.config import OpenAIModelConfig, NougatModelConfig\n",
    "from uniflow.op.prompt_schema import GuidedPrompt, Context\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb677037",
   "metadata": {},
   "source": [
    "### Prepare the input data\n",
    "First, uncomment the 10k that you want to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a707ef78",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf_file = \"nike-10k-2023.pdf\"\n",
    "# pdf_file = \"amazon-10k-2023.pdf\"\n",
    "# pdf_file = \"alphabet-10k-2023.pdf\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b177df1",
   "metadata": {},
   "source": [
    "##### Set current directory and input data directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "092b355a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_cur = os.getcwd()\n",
    "input_file = os.path.join(f\"{dir_cur}/data/raw_input/\", pdf_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load the pdf using Nougat\n",
    "For this example, we'll run the `ExtractPDF` flow to extract the text from the 10K pdf. This uses the [Nougat](https://pypi.org/project/nougat-ocr/0.1.17/) PDF parser."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = [\n",
    "    {\"pdf\": input_file},\n",
    "]\n",
    "\n",
    "extract_config = ExtractPDFConfig(\n",
    "    model_config=NougatModelConfig(\n",
    "        model_name = \"0.1.0-small\",\n",
    "        batch_size = 1 # When batch_size>1, nougat will run on CUDA, otherwise it will run on CPU\n",
    "    )\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to write a little bit prompts to generate question and answer for a given paragraph, each promopt data includes a instruction and a list of examples with \"context\", \"question\" and \"answer\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "guided_prompt = GuidedPrompt(\n",
    "    instruction=\"\"\"Generate one question and its corresponding answer based on the last context in the last\n",
    "    example. Follow the format of the examples below to include context, question, and answer in the response\"\"\",\n",
    "    examples=[\n",
    "        Context(\n",
    "            context=\"In 1948, Claude E. Shannon published A Mathematical Theory of\\nCommunication (Shannon, 1948) establishing the theory of\\ninformation. In his article, Shannon introduced the concept of\\ninformation entropy for the first time. We will begin our journey here.\",\n",
    "            question=\"Who published A Mathematical Theory of Communication in 1948?\",\n",
    "            answer=\"Claude E. Shannon.\",\n",
    "        ),\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use LLM to generate data\n",
    "\n",
    "In this example, we will use the [OpenAIModelConfig](https://github.com/CambioML/uniflow/blob/main/uniflow/model/config.py#L17)'s default LLM to generate questions and answers.\n",
    "\n",
    "Here, we pass in our `guided_prompt` to the `OpenAIConfig` to use our customized instructions and examples, instead of the `uniflow` default ones.\n",
    "\n",
    "We also want to get the response in the `json` format instead of the `text` default, so we set the `response_format` to `json_object`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform_config = TransformOpenAIConfig(\n",
    "    guided_prompt_template=guided_prompt,\n",
    "    model_config=OpenAIModelConfig(response_format={\"type\": \"json_object\"}),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we call the `run` method on the `client` object to execute the question-answer generation operation on the data shown above.\n",
    "\n",
    "Note sometimes the LLM doesn't return a JSON output, then uniflow will handle the failure and auto retry generating a new output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/cambio-recipes-10k/lib/python3.10/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3526.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n",
      "100%|██████████| 1/1 [08:31<00:00, 511.46s/it]\n",
      "100%|██████████| 1099/1099 [1:10:58<00:00,  3.88s/it]  \n"
     ]
    }
   ],
   "source": [
    "p = MultiFlowsPipeline(PipelineConfig(\n",
    "    extract_config=extract_config,\n",
    "    transform_config=transform_config,\n",
    "))\n",
    "output = p.run(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Process the output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a look of the generation output. We need to do a little postprocessing on the raw output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Context</th>\n",
       "      <th>Question</th>\n",
       "      <th>Answer</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>. Jain, Phys. Rev. Lett. **78**, 1238 (19UNITED STATES</td>\n",
       "      <td>What is the title of the article published by Jain in Phys. Rev. Lett.?</td>\n",
       "      <td>The title of the article is not provided in the given context.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SECURITIES AND EXCHANGE COMMISSION</td>\n",
       "      <td>What is the role of the SEC?</td>\n",
       "      <td>The SEC oversees and regulates the securities industry, the nation's stock and options exchanges, and other electronic securities markets.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Washington, D.C. 20549</td>\n",
       "      <td>What is the zip code for Washington, D.C.?</td>\n",
       "      <td>20549.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>FORM 10-K</td>\n",
       "      <td>What is the purpose of a FORM 10-K?</td>\n",
       "      <td>A FORM 10-K is a comprehensive report filed annually by publicly traded companies to provide a summary of their financial performance and regulatory compliance.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>FOR THE FISCAL YEAR ENDED MAY 31, 2023</td>\n",
       "      <td>What is the end date of the fiscal year mentioned?</td>\n",
       "      <td>May 31, 2023</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>913</th>\n",
       "      <td>NICE, the Swochs Design, and Just Do It are registered trademarks of NICE, Inc.</td>\n",
       "      <td>What are the registered trademarks of NICE, Inc.?</td>\n",
       "      <td>NICE, the Swochs Design, and Just Do It.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>914</th>\n",
       "      <td>NIKE, INC.</td>\n",
       "      <td>What is the name of the company?</td>\n",
       "      <td>NIKE, INC.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>915</th>\n",
       "      <td>One Bowerman Drive</td>\n",
       "      <td>What is the address of One Bowerman Drive?</td>\n",
       "      <td>One Bowerman Drive</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>916</th>\n",
       "      <td>Bawerman, OR 97005-6453</td>\n",
       "      <td>What is the zip code for Bawerman, OR?</td>\n",
       "      <td>97005-6453</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>917</th>\n",
       "      <td>www.nike.com</td>\n",
       "      <td>What is the website for Nike?</td>\n",
       "      <td>www.nike.com</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>918 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                                             Context                                                                 Question                                                                                                                                                            Answer\n",
       "0                             . Jain, Phys. Rev. Lett. **78**, 1238 (19UNITED STATES  What is the title of the article published by Jain in Phys. Rev. Lett.?                                                                                                    The title of the article is not provided in the given context.\n",
       "1                                                 SECURITIES AND EXCHANGE COMMISSION                                             What is the role of the SEC?                        The SEC oversees and regulates the securities industry, the nation's stock and options exchanges, and other electronic securities markets.\n",
       "2                                                             Washington, D.C. 20549                               What is the zip code for Washington, D.C.?                                                                                                                                                            20549.\n",
       "3                                                                          FORM 10-K                                      What is the purpose of a FORM 10-K?  A FORM 10-K is a comprehensive report filed annually by publicly traded companies to provide a summary of their financial performance and regulatory compliance.\n",
       "4                                             FOR THE FISCAL YEAR ENDED MAY 31, 2023                       What is the end date of the fiscal year mentioned?                                                                                                                                                      May 31, 2023\n",
       "..                                                                               ...                                                                      ...                                                                                                                                                               ...\n",
       "913  NICE, the Swochs Design, and Just Do It are registered trademarks of NICE, Inc.                        What are the registered trademarks of NICE, Inc.?                                                                                                                          NICE, the Swochs Design, and Just Do It.\n",
       "914                                                                       NIKE, INC.                                         What is the name of the company?                                                                                                                                                        NIKE, INC.\n",
       "915                                                               One Bowerman Drive                               What is the address of One Bowerman Drive?                                                                                                                                                One Bowerman Drive\n",
       "916                                                          Bawerman, OR 97005-6453                                   What is the zip code for Bawerman, OR?                                                                                                                                                        97005-6453\n",
       "917                                                                     www.nike.com                                            What is the website for Nike?                                                                                                                                                      www.nike.com\n",
       "\n",
       "[918 rows x 3 columns]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Extracting context, question, and answer into a DataFrame\n",
    "contexts = []\n",
    "questions = []\n",
    "answers = []\n",
    "\n",
    "for item in output[0]:\n",
    "    for i in item.get('output', []):\n",
    "        for response in i.get('response', []):\n",
    "            if any(key not in response for key in ['context', 'question', 'answer']):\n",
    "                print(\"[WARNING] Missing context, question or answer in response, skipping:\\n\", response)\n",
    "                continue\n",
    "            if \"Claude E. Shannon\" in response['context']:\n",
    "                print(\"[WARNING] Used example context, skipping:\\n\", response[\"context\"])\n",
    "                continue\n",
    "            contexts.append(response['context'])\n",
    "            questions.append(response['question'])\n",
    "            answers.append(response['answer'])\n",
    "\n",
    "# Set display options\n",
    "pd.set_option('display.max_colwidth', None)  # or use a specific width like 50\n",
    "pd.set_option('display.width', 1000)\n",
    "\n",
    "df = pd.DataFrame({\n",
    "    'Context': contexts,\n",
    "    'Question': questions,\n",
    "    'Answer': answers\n",
    "})\n",
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can save the `uniflow` output to a `.csv` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "72a570e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_df = df[['Question', 'Answer']]\n",
    "\n",
    "output_dir = 'data/output'\n",
    "\n",
    "uniflow_output_path = f\"{output_dir}/Nike_10k_QApairs.csv\"\n",
    "\n",
    "if not os.path.exists(output_dir):\n",
    "    os.makedirs(output_dir)\n",
    "\n",
    "output_df.to_csv(uniflow_output_path, index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Release GPU Memory\n",
    "We'll need to use our GPU for future steps, so let's release the memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU memory has been released.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.empty_cache()\n",
    "    print(\"GPU memory has been released.\")\n",
    "else:\n",
    "    print(\"No GPU devices found.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Running `pykoi` `SupervisedFineTuning` on the QA pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install helper packages\n",
    "If you already have these installed, feel free to skip this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: ...\n"
     ]
    }
   ],
   "source": [
    "!{sys.executable} -m pip install peft"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Dependency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/cambio-recipes-10k/lib/python3.10/site-packages/trl/trainer/ppo_config.py:141: UserWarning: The `optimize_cuda_cache` arguement will be deprecated soon, please use `optimize_device_cache` instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from pykoi.rlhf import RLHFConfig\n",
    "from pykoi.rlhf import SupervisedFinetuning\n",
    "from peft import LoraConfig, TaskType"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set the parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_model_path = \"meta-llama/Llama-2-7b-chat-hf\"\n",
    "dataset_name = uniflow_output_path\n",
    "peft_model_path = \"./models/rlhf_step1_sft\"\n",
    "dataset_type = \"local_csv\"\n",
    "learning_rate = 1e-3\n",
    "weight_decay = 0.0\n",
    "max_steps = 1600\n",
    "per_device_train_batch_size = 1\n",
    "per_device_eval_batch_size = 4\n",
    "log_freq = 20\n",
    "eval_freq = 2000\n",
    "save_freq = 200\n",
    "train_test_split_ratio = 0.0001\n",
    "dataset_subset_sft_train = 999999999\n",
    "size_valid_set = 0\n",
    "\n",
    "r = 8\n",
    "lora_alpha = 16\n",
    "lora_dropout = 0.05\n",
    "bias = \"none\"\n",
    "task_type = TaskType.CAUSAL_LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "lora_config = LoraConfig(\n",
    "    r=r,\n",
    "    lora_alpha=lora_alpha,\n",
    "    lora_dropout=lora_dropout,\n",
    "    bias=bias,\n",
    "    task_type=task_type,\n",
    "    )\n",
    "\n",
    "\n",
    "# run supervised finetuning\n",
    "config = RLHFConfig(\n",
    "    base_model_path=base_model_path,\n",
    "    dataset_type=dataset_type,\n",
    "    dataset_name=dataset_name,\n",
    "    learning_rate=learning_rate,\n",
    "    weight_decay=weight_decay,\n",
    "    max_steps=max_steps,\n",
    "    per_device_train_batch_size=per_device_train_batch_size,\n",
    "    per_device_eval_batch_size=per_device_eval_batch_size,\n",
    "    log_freq=log_freq,\n",
    "    eval_freq=eval_freq,\n",
    "    save_freq=save_freq,\n",
    "    train_test_split_ratio=train_test_split_ratio,\n",
    "    dataset_subset_sft_train=dataset_subset_sft_train,\n",
    "    size_valid_set=size_valid_set,\n",
    "    lora_config_rl=lora_config\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the SupervisedFineTuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading data files: 100%|██████████| 1/1 [00:00<00:00, 10407.70it/s]\n",
      "Extracting data files: 100%|██████████| 1/1 [00:00<00:00, 1367.11it/s]\n",
      "Generating train split: 918 examples [00:00, 110585.65 examples/s]\n",
      "/opt/conda/envs/cambio-recipes-10k/lib/python3.10/site-packages/trl/trainer/utils.py:548: UserWarning: The passed formatting_func has more than one argument. Usually that function should have a single argument `example` which corresponds to the dictionary returned by each element of the dataset. Make sure you know what you are doing.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of the train set: 917.               Size of the validation set: 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 2/2 [00:01<00:00,  1.03it/s]\n",
      "/opt/conda/envs/cambio-recipes-10k/lib/python3.10/site-packages/trl/trainer/sft_trainer.py:194: UserWarning: You didn't pass a `max_seq_length` argument to the SFTTrainer, this will default to 1024\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/cambio-recipes-10k/lib/python3.10/site-packages/trl/trainer/sft_trainer.py:267: UserWarning: You passed `packing=True` to the SFTTrainer, and you are training your model with `max_steps` strategy. The dataset will be iterated until the `max_steps` are reached.\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/cambio-recipes-10k/lib/python3.10/site-packages/trl/trainer/utils.py:570: UserWarning: The dataset reached end and the iterator is reset to the start.\n",
      "  warnings.warn(\"The dataset reached end and the iterator is reset to the start.\")\n",
      "You're using a LlamaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n",
      "/opt/conda/envs/cambio-recipes-10k/lib/python3.10/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1600' max='1600' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1600/1600 1:42:42, Epoch 6/7]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "rlhf_step1_sft = SupervisedFinetuning(config)\n",
    "rlhf_step1_sft.train_and_save(peft_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Release GPU Memory\n",
    "We'll need to use our GPU for future steps, so let's release the memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU memory has been released.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.empty_cache()\n",
    "    print(\"GPU memory has been released.\")\n",
    "else:\n",
    "    print(\"No GPU devices found.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Running a `pykoi` `Chatbot` on the fine-tuned model\n",
    "\n",
    "### Import pykoi components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pykoi.application import Application\n",
    "from pykoi.chat import ModelFactory\n",
    "from pykoi.chat import QuestionAnswerDatabase\n",
    "from pykoi.component import Chatbot, Dashboard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/cambio-recipes-10k/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HuggingfaceModel] loading base model...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 2/2 [01:42<00:00, 51.00s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HuggingfaceModel] loading perf model...\n",
      "[HuggingfaceModel] loading tokenizer...\n"
     ]
    }
   ],
   "source": [
    "model = ModelFactory.create_model(\n",
    "    model_source=\"peft_huggingface\",\n",
    "    base_model_path=\"meta-llama/Llama-2-7b-chat-hf\",\n",
    "    lora_model_path=\"/home/ubuntu/pykoi/models/rlhf_step1_sft\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the Chatbot with the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Table contents after creating table:\n",
      "ID: 1, Question: Who is on Nike's board, Answer: Who is on Nike's board of directors?[Page 145]\n",
      "\n",
      "            Answer: John J. Donahoe II, Matthew Friend, Johanna Nielsen, Mark G. Parker, Cathleen A. Benko, Timothy D. Cook, Thasunda B. Duckett, Mónica Gil, Alan B. Graf, Jr., Maria Henry, Peter B. Henry, Travis A. Knight, Michelle A., Vote Status: n/a, Timestamp: 2023-12-28 17:22:18.342838\n"
     ]
    }
   ],
   "source": [
    "database = QuestionAnswerDatabase(debug=True)\n",
    "chatbot = Chatbot(model=model, feedback=\"vote\")\n",
    "dashboard = Dashboard(database=database)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the Chatbot app!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Add `nest_asyncio` \n",
    "Add `nest_asyncio` to avoid error such as `asyncio.run() cannot be called from a running event loop`. Since we're running another interface inside a Jupyter notebook where an asyncio event loop is already running, we'll encounter the error. (since The uvicorn.run() function uses asyncio.run(), which isn't compatible with a running event loop.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -q nest_asyncio\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:     Started server process [2254]\n",
      "INFO:     Waiting for application startup.\n",
      "INFO:     Application startup complete.\n",
      "INFO:     Uvicorn running on http://0.0.0.0:5000 (Press CTRL+C to quit)\n"
     ]
    }
   ],
   "source": [
    "app = Application(debug=False, share=False)\n",
    "app.add_component(chatbot)\n",
    "app.add_component(dashboard)\n",
    "app.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## End of the notebook\n",
    "\n",
    "Check more use cases in the [example folder](../../examples/)!\n",
    "\n",
    "<a href=\"https://www.cambioml.com/\" title=\"Title\">\n",
    "    <img src=\"../image/cambioml_logo_large.png\" style=\"height: 100px; display: block; margin-left: auto; margin-right: auto;\"/>\n",
    "</a>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "self-instruct-ft",
   "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.1.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}