02. Generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/wandb/edu/blob/main/llm-apps-course/notebooks/02.%20Generation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "<!--- @wandbcode{llmapps-generation} -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generation\n",
    "<!--- @wandbcode{llmapps-generation} -->\n",
    "\n",
    "In this notebook we will dive deeper on prompting the model by passing a better context by using available data from users questions and using the documentation files to generate better answers.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -Uqqq rich openai==0.27.2 tiktoken wandb tenacity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "\n",
    "import openai\n",
    "import tiktoken\n",
    "\n",
    "from pathlib import Path\n",
    "from pprint import pprint\n",
    "from getpass import getpass\n",
    "\n",
    "from rich.markdown import Markdown\n",
    "import pandas as pd\n",
    "from tenacity import (\n",
    "    retry,\n",
    "    stop_after_attempt,\n",
    "    wait_random_exponential, # for exponential backoff\n",
    ")  \n",
    "import wandb\n",
    "from wandb.integration.openai import autolog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download files on colab\n",
    "if not Path(\"examples.txt\").exists():\n",
    "    !wget https://raw.githubusercontent.com/wandb/edu/main/llm-apps-course/notebooks/{examples,prompt_template,system_template}.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will need an OpenAI API key to run this notebook. You can get one [here](https://platform.openai.com/account/api-keys)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.getenv(\"OPENAI_API_KEY\") is None:\n",
    "  if any(['VSCODE' in x for x in os.environ.keys()]):\n",
    "    print('Please enter password in the VS Code prompt at the top of your VS Code window!')\n",
    "  os.environ[\"OPENAI_API_KEY\"] = getpass(\"Paste your OpenAI key from: https://platform.openai.com/account/api-keys\\n\")\n",
    "  openai.api_key = os.getenv(\"OPENAI_API_KEY\", \"\")\n",
    "\n",
    "assert os.getenv(\"OPENAI_API_KEY\", \"\").startswith(\"sk-\"), \"This doesn't look like a valid OpenAI API key\"\n",
    "print(\"OpenAI API key configured\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's enable W&B autologging to track our experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start logging to W&B\n",
    "autolog({\"project\":\"llmapps\", \"job_type\": \"generation\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating synthetic support questions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will add a retry behavior in case we hit the API rate limit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))\n",
    "def completion_with_backoff(**kwargs):\n",
    "    return openai.ChatCompletion.create(**kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"gpt-3.5-turbo\"\n",
    "# MODEL_NAME = \"gpt-4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"You are a helpful assistant.\"\n",
    "user_prompt = \"Generate a support question from a W&B user\"\n",
    "\n",
    "def generate_and_print(system_prompt, user_prompt, n=5):\n",
    "    messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt},\n",
    "        ]\n",
    "    responses = completion_with_backoff(\n",
    "        model=MODEL_NAME,\n",
    "        messages=messages,\n",
    "        n = n,\n",
    "        )\n",
    "    for response in responses.choices:\n",
    "        generation = response.message.content\n",
    "        display(Markdown(generation))\n",
    "    \n",
    "generate_and_print(system_prompt, user_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Few Shot "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's read some user submitted queries from the file `examples.txt`. This file contains multiline questions separated by tabs (`\\t`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test if examples.txt is present, download if not\n",
    "if not Path(\"examples.txt\").exists():\n",
    "    !wget https://raw.githubusercontent.com/wandb/edu/main/llm-apps-course/notebooks/examples.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delimiter = \"\\t\" # tab separated queries\n",
    "with open(\"examples.txt\", \"r\") as file:\n",
    "    data = file.read()\n",
    "    real_queries = data.split(delimiter)\n",
    "\n",
    "pprint(f\"We have {len(real_queries)} real queries:\")  \n",
    "Markdown(f\"Sample one: \\n\\\"{random.choice(real_queries)}\\\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now use those real user questions to guide our model to produce synthetic questions like those."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_few_shot_prompt(queries, n=3):\n",
    "    prompt = \"Generate a support question from a W&B user\\n\" +\\\n",
    "        \"Below you will find a few examples of real user queries:\\n\"\n",
    "    for _ in range(n):\n",
    "        prompt += random.choice(queries) + \"\\n\"\n",
    "    prompt += \"Let's start!\"\n",
    "    return prompt\n",
    "\n",
    "generation_prompt = generate_few_shot_prompt(real_queries)\n",
    "Markdown(generation_prompt)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OpenAI `Chat` models are really good at following instructions with a few examples. Let's see how it does here. This is going to use some context from the prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_and_print(system_prompt, user_prompt=generation_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Add Context & Response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a function to find all the markdown files in a directory and return it's content and path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if directory exists, if not, create it and download the files, e.g if running in colab\n",
    "if not os.path.exists(\"../docs_sample/\"):\n",
    "  !git clone https://github.com/wandb/edu.git\n",
    "  !cp -r edu/llm-apps-course/docs_sample ../"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_md_files(directory):\n",
    "    \"Find all markdown files in a directory and return their content and path\"\n",
    "    md_files = []\n",
    "    for file in Path(directory).rglob(\"*.md\"):\n",
    "        with open(file, 'r', encoding='utf-8') as md_file:\n",
    "            content = md_file.read()\n",
    "        md_files.append((file.relative_to(directory), content))\n",
    "    return md_files\n",
    "\n",
    "documents = find_md_files('../docs_sample/')\n",
    "len(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check if the documents are not too long for our context window. We need to compute the number of tokens in each document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = tiktoken.encoding_for_model(MODEL_NAME)\n",
    "tokens_per_document = [len(tokenizer.encode(document)) for _, document in documents]\n",
    "pprint(tokens_per_document)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some of them are too long - instead of using entire documents, we'll extract a random chunk from them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract a random chunk from a document\n",
    "def extract_random_chunk(document, max_tokens=512):\n",
    "    tokens = tokenizer.encode(document)\n",
    "    if len(tokens) <= max_tokens:\n",
    "        return document\n",
    "    start = random.randint(0, len(tokens) - max_tokens)\n",
    "    end = start + max_tokens\n",
    "    return tokenizer.decode(tokens[start:end])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we will use that extracted chunk to create a question that can be answered by the document. This way we can generate questions that our current documentation is capable of answering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_context_prompt(chunk):\n",
    "    prompt = \"Generate a support question from a W&B user\\n\" +\\\n",
    "        \"The question should be answerable by provided fragment of W&B documentation.\\n\" +\\\n",
    "        \"Below you will find a fragment of W&B documentation:\\n\" +\\\n",
    "        chunk + \"\\n\" +\\\n",
    "        \"Let's start!\"\n",
    "    return prompt\n",
    "\n",
    "chunk = extract_random_chunk(documents[0][1])\n",
    "generation_prompt = generate_context_prompt(chunk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Markdown(generation_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's generate 3 possible questions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_and_print(system_prompt, generation_prompt, n=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> As you can see, sometimes the generation contains an intro phrase like: \"Sure, here's a support question based on the documentation:\", we may want to put some instructions to avoid this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Level 5 prompt\n",
    "\n",
    "Complex directive that includes the following:\n",
    "- Description of high-level goal\n",
    "- A detailed bulleted list of sub-tasks\n",
    "- An explicit statement asking LLM to explain its own output\n",
    "- A guideline on how LLM output will be evaluated\n",
    "- Few-shot examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we will use GPT4 from here, as it gives better answers and abides to instructions better\n",
    "MODEL_NAME = \"gpt-4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read system_template.txt file into an f-string\n",
    "with open(\"system_template.txt\", \"r\") as file:\n",
    "    system_prompt = file.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Markdown(system_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read prompt_template.txt file into an f-string\n",
    "with open(\"prompt_template.txt\", \"r\") as file:\n",
    "    prompt_template = file.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Markdown(prompt_template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_context_prompt(chunk, n_questions=3):\n",
    "    questions = '\\n'.join(random.sample(real_queries, n_questions))\n",
    "    user_prompt = prompt_template.format(QUESTIONS=questions, CHUNK=chunk)\n",
    "    return user_prompt\n",
    "\n",
    "user_prompt = generate_context_prompt(chunk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Markdown(user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_questions(documents, n_questions=3, n_generations=5):\n",
    "    questions = []\n",
    "    for _, document in documents:\n",
    "        chunk = extract_random_chunk(document)\n",
    "        user_prompt = generate_context_prompt(chunk, n_questions)\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt},\n",
    "        ]\n",
    "        response = completion_with_backoff(\n",
    "            model=MODEL_NAME,\n",
    "            messages=messages,\n",
    "            n = n_generations,\n",
    "            )\n",
    "        questions.extend([response.choices[i].message.content for i in range(n_generations)])\n",
    "    return questions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> A Note about the `system` role: For GPT4 based pipelines you probably want to move some part of the context prompt to the `system` context. As we are using `gpt3.5-turbo` here, you can put the instruction on the user prompt, you can read more about this on [OpenAI docs here](https://platform.openai.com/docs/guides/chat/instructing-chat-models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to parse model generation and extract CONTEXT, QUESTION and ANSWER\n",
    "def parse_generation(generation):\n",
    "    lines = generation.split(\"\\n\")\n",
    "    context = []\n",
    "    question = []\n",
    "    answer = []\n",
    "    flag = None\n",
    "    \n",
    "    for line in lines:\n",
    "        if \"CONTEXT:\" in line:\n",
    "            flag = \"context\"\n",
    "            line = line.replace(\"CONTEXT:\", \"\").strip()\n",
    "        elif \"QUESTION:\" in line:\n",
    "            flag = \"question\"\n",
    "            line = line.replace(\"QUESTION:\", \"\").strip()\n",
    "        elif \"ANSWER:\" in line:\n",
    "            flag = \"answer\"\n",
    "            line = line.replace(\"ANSWER:\", \"\").strip()\n",
    "\n",
    "        if flag == \"context\":\n",
    "            context.append(line)\n",
    "        elif flag == \"question\":\n",
    "            question.append(line)\n",
    "        elif flag == \"answer\":\n",
    "            answer.append(line)\n",
    "\n",
    "    context = \"\\n\".join(context)\n",
    "    question = \"\\n\".join(question)\n",
    "    answer = \"\\n\".join(answer)\n",
    "    return context, question, answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generations = generate_questions([documents[0]], n_questions=3, n_generations=5)\n",
    "parse_generation(generations[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parsed_generations = []\n",
    "generations = generate_questions(documents, n_questions=3, n_generations=5)\n",
    "for generation in generations:\n",
    "    context, question, answer = parse_generation(generation)\n",
    "    parsed_generations.append({\"context\": context, \"question\": question, \"answer\": answer})\n",
    "\n",
    "# let's convert parsed_generations to a pandas dataframe and save it locally\n",
    "df = pd.DataFrame(parsed_generations)\n",
    "df.to_csv('generated_examples.csv', index=False)\n",
    "\n",
    "# log df as a table to W&B for interactive exploration\n",
    "wandb.log({\"generated_examples\": wandb.Table(dataframe=df)})\n",
    "\n",
    "# log csv file as an artifact to W&B for later use\n",
    "artifact = wandb.Artifact(\"generated_examples\", type=\"dataset\")\n",
    "artifact.add_file(\"generated_examples.csv\")\n",
    "wandb.log_artifact(artifact)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}