gpt4-rag-logprobs/gpt4-rag-logprobs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Obtaining logprobs from GPT-4 RAG systems\n",
    "\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook showcases how the GPT-4 logprob results were obtained in our tutorial on integrating Cleanlab TLM within a RAG system. We reference code from OpenAI's [blogpost](https://cookbook.openai.com/examples/using_logprobs) on logprobs.\n",
    "\n",
    "From the blogpost,\n",
    "> Log probabilities of output tokens indicate the likelihood of each token occurring in the sequence given the context. To simplify, a logprob is log(p), where p = probability of a token occurring at a specific position based on the previous tokens in the context."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installing LlamaIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize OpenAI Client\n",
    "\n",
    "LlamaIndex uses OpenAI’s embedding models by default. Make sure your API key is activated in your environment by using the following command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "import os\n",
    "API_KEY = \"<API Key>\"\n",
    "os.environ['OPENAI_API_KEY'] = API_KEY\n",
    "client = OpenAI(api_key=os.environ.get(\"OPENAI_API_KEY\", API_KEY))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Integrating GPT-4 with LlamaIndex"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code defines two utility functions. `get_completion()` returns the output from the OpenAI client given a prompt and a few other customizable parameters. `parse()` extracts our desired fields from the output, `response` and `logprobs`, and calculates the average logprob and linear probability over token in the `response` string."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import exp\n",
    "import numpy as np\n",
    "from IPython.display import display, HTML\n",
    "from typing import Dict, List, Optional\n",
    "\n",
    "def get_completion(\n",
    "    messages: List[Dict[str, str]],\n",
    "    model: str = \"gpt-4\",\n",
    "    max_tokens=500,\n",
    "    temperature=0,\n",
    "    stop=None,\n",
    "    seed=123,\n",
    "    tools=None,\n",
    "    logprobs=None,  # whether to return log probabilities of the output tokens or not. If true, returns the log probabilities of each output token returned in the content of message..\n",
    "    top_logprobs=None,\n",
    ") -> str:\n",
    "    params = {\n",
    "        \"model\": model,\n",
    "        \"messages\": messages,\n",
    "        \"max_tokens\": max_tokens,\n",
    "        \"temperature\": temperature,\n",
    "        \"stop\": stop,\n",
    "        \"seed\": seed,\n",
    "        \"logprobs\": logprobs,\n",
    "        \"top_logprobs\": top_logprobs,\n",
    "    }\n",
    "    if tools:\n",
    "        params[\"tools\"] = tools\n",
    "\n",
    "    completion = client.chat.completions.create(**params)\n",
    "    return completion\n",
    "\n",
    "def parse(api_response):\n",
    "    choice = api_response.choices[0]\n",
    "    content = choice.message.content\n",
    "    logprobs = [logprob.logprob for logprob in choice.logprobs.content]\n",
    "    \n",
    "    average_logprob = np.mean(logprobs)\n",
    "    average_linear_prob = np.exp(average_logprob) * 100\n",
    "    \n",
    "    result_string = (f\"Response: {content}\\n\"\n",
    "                     f\"Average Log Probability: {average_logprob:.4f}\\n\"\n",
    "                     f\"Average Linear Probability: {average_linear_prob:.2f}%\")\n",
    "    \n",
    "    return result_string"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code sets GPT-4 (with logprobs) as the underlying large language model (LLM) for our RAG system. `GPTWrapper` is built on top of LlamaIndex's [CustomLLM](https://docs.llamaindex.ai/en/stable/module_guides/models/llms/usage_custom.html#using-custom-llm-advanced) class. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Callable, Optional, Sequence\n",
    "\n",
    "from llama_index.core.base.llms.types import (\n",
    "    ChatMessage,\n",
    "    CompletionResponse,\n",
    "    CompletionResponseGen,\n",
    "    LLMMetadata,\n",
    ")\n",
    "from llama_index.core.callbacks import CallbackManager\n",
    "from llama_index.core.llms.callbacks import llm_completion_callback\n",
    "from llama_index.core.llms.custom import CustomLLM\n",
    "from llama_index.core.types import PydanticProgramMode\n",
    "from llama_index.core import Settings\n",
    "import json\n",
    "\n",
    "class GPTWrapper(CustomLLM):\n",
    "    context_window: int = 3900\n",
    "    num_output: int = 256\n",
    "    model_name: str = \"ChatGPT-4\"\n",
    "\n",
    "    @property\n",
    "    def metadata(self) -> LLMMetadata:\n",
    "        \"\"\"Get LLM metadata.\"\"\"\n",
    "        return LLMMetadata(\n",
    "            context_window=self.context_window,\n",
    "            num_output=self.num_output,\n",
    "            model_name=self.model_name,\n",
    "        )\n",
    "\n",
    "    @llm_completion_callback()\n",
    "    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:\n",
    "        API_RESPONSE = get_completion(\n",
    "            [\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": prompt\n",
    "                }\n",
    "            ],\n",
    "            model=\"gpt-4\",\n",
    "            logprobs=True,\n",
    "        )\n",
    "        return CompletionResponse(text=parse(API_RESPONSE))\n",
    "\n",
    "\n",
    "    @llm_completion_callback()\n",
    "    def stream_complete(\n",
    "        self, prompt: str, **kwargs: Any\n",
    "    ) -> CompletionResponseGen:\n",
    "        API_RESPONSE = get_completion(\n",
    "            [\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": prompt\n",
    "                }\n",
    "            ],\n",
    "            model=\"gpt-4\",\n",
    "            logprobs=True,\n",
    "        )\n",
    "        for char in API_RESPONSE:\n",
    "            yield CompletionResponse(text=char, delta=char)\n",
    "\n",
    "Settings.llm = GPTWrapper()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data and build an index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial uses Nvidia's earnings report from Q1 FY2024. Download it via the command below and save it in a folder called `data`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wget -nc 'https://cleanlab-public.s3.amazonaws.com/Datasets/NVIDIA_Financial_Results_Q1_FY2024.md'\n",
    "!mkdir -p data\n",
    "!mv NVIDIA_Financial_Results_Q1_FY2024.md data/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your directory structure should look like this:\n",
    "```\n",
    "├── tlm-rag-tutorial.ipynb\n",
    "└── data\n",
    "    └── NVIDIA_Financial_Results_Q1_FY2024.md\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell constructs an index from the content in the `data` folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\"data\").load_data()\n",
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Query your data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, you can create an Q&A engine over your index and input prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response: The document does not provide specific information on what Nvidia's largest product is.\n",
      "Average Log Probability: -0.0582\n",
      "Average Linear Probability: 94.34%\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"What is Nvidia's largest product?\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response: False\n",
      "Average Log Probability: -0.0081\n",
      "Average Linear Probability: 99.19%\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"True or False: Nvidia's Professional Visualization division is performing better than their Gaming division in terms of percent change in revenue compared to the previous quarter.\")\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ag",
   "language": "python",
   "name": "ag"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}