synthetic_preference_data_generation_nemotron_4_340B.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Synthetic Preference Data Generation Using Nemotron-4 340B\n",
    "\n",
    "The following notebook will demonstrate how to leverage [Nemotron-4 340B Instruct](https://build.nvidia.com/nvidia/nemotron-4-340b-instruct), and [Nemotron-4 340B Reward](https://build.nvidia.com/nvidia/nemotron-4-340b-reward) through [build.nvidia.com](https://build.nvidia.com/explore/discover).\n",
    "\n",
    "The build will be a demonstration of the following pipeline, as discuss in the [release blog](https://blogs.nvidia.com/blog/nemotron-4-synthetic-data-generation-llm-training/), and [technical blog](https://developer.nvidia.com/blog/leverage-our-latest-open-models-for-synthetic-data-generation-with-nvidia-nemotron-4-340b/). The pipeline is designed to create a preference dataset suitable for training a custom reward model using the [SteerLM method](https://docs.nvidia.com/nemo-framework/user-guide/latest/modelalignment/steerlm.html), however consecutive responses (e.g. sample 1 with 2, 3 with 4, etc.) share the same prompt so the dataset can also be used for preference pairs for training an RLHF Reward Model or for DPO - using the helpfulness score.\n",
    "\n",
    "![image](https://developer-blogs.nvidia.com/wp-content/uploads/2024/06/SDG-Pipeline-1-625x352.png)\n",
    "\n",
    "The flow will be split into 2 general parts: \n",
    "\n",
    "1. **Synthetic Response Generation**: A domain specific input query will be provided by the developer - at which point Nemotron-4 340B Instruct will be leveraged to generate ~150 questions. Then, Nemotron-4 340B Instruct will be used to generated 2 responses for each question. \n",
    "2. **Reward Model as a Judge**: Nemotron-4 340B Reward will be used to score the 2 responses per question to be used for further alignment training via [NeMo Aligner](https://github.com/NVIDIA/NeMo-Aligner)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## build.nvidia.com API Key Set-up!\n",
    "\n",
    "In order to access the endpoints through [build.nvidia.com](https://build.nvidia.com/explore/discover), an API key is required. \n",
    "\n",
    "A trial API key is made available with 1,000 tokens (or 5,000 tokens for corporate emails) - the example below will leverage ~4,500 tokens of data, but can be extended beyond that limit using local instances of the models.\n",
    "\n",
    "There are two steps to get a trial API key:\n",
    "\n",
    "1. Login (or sign up) through [build.nvidia.com](https://build.nvidia.com/)\n",
    "2. Click the `Get API Key` button available on the the `nvidia/nemotron-4-340b-instruct` page, found [here](https://build.nvidia.com/nvidia/nemotron-4-340b-instruct).\n",
    "\n",
    "![image](https://i.imgur.com/dM7AwKZ.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Generate Subtopics, questions, and responses with Nemotron-4 340B Instruct\n",
    "\n",
    "The first part of the notebook will cover the creation of raw synthetic data from Nemotron-4 340B Instruct, due to the model's [permissive license](https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf), the usage of the outputs of Nemotron-4 340B Instruct are permitted to be used for training, customization, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prompt Templates for Synthetic Data Generation\n",
    "\n",
    "To generate questions and responses, there are a few prompt templates required:\n",
    "\n",
    "1. A prompt template to generate subtopics from a user provided topic\n",
    "2. A prompt template to generate questions for a given subtopic\n",
    "2. A prompt template to generate responses for a given question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "metadata": {},
   "outputs": [],
   "source": [
    "TOPIC_GENERATION_PROMPT_TEMPLATE = \"\"\"\\\n",
    "Given a topic, generate a list of {n_subtopics} subtopics that are related to the topic.\n",
    "\n",
    "The topic is: {topic}\n",
    "\n",
    "The list must be without numbers, and without any description of the subtopics. The subtopics should be separated by a comma. There must be no other text than the list.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "metadata": {},
   "outputs": [],
   "source": [
    "QUESTION_PROMPT_TEMPLATE = \"\"\"\\\n",
    "Given a topic, generate {n_questions} questions that could be asked about that topic. Your response should be in a list format.\n",
    "\n",
    "The topic is: {sub_topic}\n",
    "\n",
    "The list must be without numbers. The questions should be separated by a newline character. There must be no other text than the list.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {},
   "outputs": [],
   "source": [
    "RESPONSE_PROMPT_TEMPLATE = \"\"\"\\\n",
    "Given a question, generate 2 responses that could be given to that question. Your response should be in a list format.\n",
    "\n",
    "The question is: {question}\n",
    "\n",
    "The list must be in the format:\n",
    "\n",
    "RESPONSE A: Response A text here\n",
    "RESPONSE B: Response B text here\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defined below are the parameters that will be used throughout the notebook to generate numbers of datapoints. \n",
    "\n",
    "1. `n_subtopics`, for the given topic `10` sub-topics will be generated by Nemotron-4 340B Instruct\n",
    "2. `n_questions`, for the given sub-topic, `10` questions will be generated by Nemotron-4 340B Instruct\n",
    "\n",
    "> NOTE: Using the default parameters above - there will be 10 sub-topics, each with 10 questions, each with 2 (hardcoded) responses. That is a total of an estimated ~200 rows of data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_subtopics = 10\n",
    "n_questions = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting OpenAI Client for Synthetic Data Generation\n",
    "\n",
    "Due to [build.nvidia.com](https://build.nvidia.com/)'s integration with the OpenAI API template - the OpenAI Python library can be used to interact with Nemotron-4 340B Instruct and Nemotron-4 340B Reward.\n",
    "\n",
    "To begin, install the [OpenAI Python library](https://github.com/openai/openai-python)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qU openai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Provide the NVIDIA API key obtained above in order to ensure access to both models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import getpass\n",
    "\n",
    "os.environ[\"NVIDIA_API_KEY\"] = getpass.getpass(\"Please enter your OpenAI API key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the OpenAI Async client will enable quick and efficient data generation.\n",
    "\n",
    "It's as easy as pointing the `base_url` parameter to `https://integrate.api.nvidia.com/v1` - and providing the API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import AsyncOpenAI\n",
    "\n",
    "client = AsyncOpenAI(\n",
    "  base_url = \"https://integrate.api.nvidia.com/v1\",\n",
    "  api_key = os.environ[\"NVIDIA_API_KEY\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating Subtopics\n",
    "\n",
    "To start things off, subtopics will be generated for the provided topic. \n",
    "\n",
    "> NOTE: The parameters of `temperature`, `top_p`, and `max_tokens` can be customized to individual preference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def generate_subtopics(client, topic, n_subtopics):\n",
    "    prompt = TOPIC_GENERATION_PROMPT_TEMPLATE.format(topic=topic, n_subtopics=n_subtopics)\n",
    "    response = await client.chat.completions.create(\n",
    "        model=\"nvidia/nemotron-4-340b-instruct\",\n",
    "        messages=[\n",
    "            {\"role\" : \"user\",\n",
    "             \"content\" : prompt}\n",
    "        ],\n",
    "        temperature=0.2,\n",
    "        top_p=0.7,\n",
    "        max_tokens=1024,\n",
    "    )\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main topic can be defined below - for the example in the notebook, \"Machine Learning\" will be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [],
   "source": [
    "topic = \"Machine Learning\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below will call the Nemotron-4 340B Instruct endpoint - and return a list of subtopics separated by commas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [],
   "source": [
    "responses = await generate_subtopics(client, topic=topic, n_subtopics=n_subtopics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output conforms to the expected format below.\n",
    "\n",
    "> NOTE: It is possible that additional data cleaning, or formatting may be necessary depending on the prompt templates used. Be sure to confirm the format of the generated data at each step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Supervised learning, unsupervised learning, semi-supervised learning, reinforcement learning, deep learning, neural networks, natural language processing, computer vision, recommendation systems, anomaly detection.\n"
     ]
    }
   ],
   "source": [
    "print(responses.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Due to the data being generated in a comma separated list, Python's `.split(\",\")` will convert the string into a usable list for the following steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [],
   "source": [
    "subtopic_list = responses.choices[0].message.content.split(\",\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating Questions from Subtopic List\n",
    "\n",
    "With a list of subtopics, the next step will be to generate `n_questions`, for each subtopic.\n",
    "\n",
    "First, there needs to be a function to generate \"batches\" of questions.\n",
    "\n",
    "> NOTE: It would suitable to generate a single question per topic at a time, but more care would be needed to confirm there were no duplicate questions in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def generate_questions(client, sub_topic, n_questions):\n",
    "    prompt = QUESTION_PROMPT_TEMPLATE.format(sub_topic=sub_topic, n_questions=n_questions)\n",
    "    response = await client.chat.completions.create(\n",
    "        model=\"nvidia/nemotron-4-340b-instruct\",\n",
    "        messages=[\n",
    "            {\"role\" : \"user\",\n",
    "             \"content\" : prompt}\n",
    "        ],\n",
    "        temperature=0.2,\n",
    "        top_p=0.7,\n",
    "        max_tokens=1024,\n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This step leverages [`asyncio`](https://docs.python.org/3/library/asyncio.html) from Python's standard library for efficient API calls to [build.nvidia.com](https://build.nvidia.com/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "\n",
    "async def question_generator(client, subtopic_list, n_question):\n",
    "    tasks = [generate_questions(client, subtopic, n_question) for subtopic in subtopic_list]\n",
    "    question_list = await asyncio.gather(*tasks)\n",
    "    return question_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Due to running in a Colab environment - it is necessary to use `nest_asyncio` to run an event loop during the current Jupyter event loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()\n",
    "\n",
    "question_list = asyncio.run(question_generator(client, subtopic_list, n_questions))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's time to examine the output of the above process!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['What is supervised learning and how does it differ from unsupervised learning?\\nWhat are the key components of a supervised learning model?\\nCan you explain the concept of labeled data in supervised learning?\\nWhat are some popular algorithms used in supervised learning?\\nHow do you evaluate the performance of a supervised learning model?\\nWhat is overfitting and how can it be prevented in supervised learning?\\nHow do you choose the right supervised learning algorithm for a given problem?\\nWhat is the role of feature engineering in supervised learning?\\nCan you explain the concept of cross-validation in supervised learning?\\nWhat are some real-world applications of supervised learning?',\n",
       " 'What is unsupervised learning and how does it differ from supervised learning?\\nWhat are some common algorithms used in unsupervised learning?\\nHow can unsupervised learning be used for anomaly detection?\\nWhat is clustering and what are some popular clustering algorithms?\\nHow can dimensionality reduction be achieved through unsupervised learning?\\nWhat is the difference between principal component analysis and t-distributed stochastic neighbor embedding?\\nHow can unsupervised learning be used for natural language processing tasks?\\nWhat is the role of unsupervised learning in deep learning?\\nHow can the performance of unsupervised learning algorithms be evaluated?\\nWhat are some real-world applications of unsupervised learning?',\n",
       " 'What is semi-supervised learning and how does it differ from supervised and unsupervised learning?\\nHow can semi-supervised learning be used to improve the performance of machine learning models?\\nWhat are some common techniques used in semi-supervised learning, such as self-training and multi-view training?\\nHow can semi-supervised learning be applied to image and text classification tasks?\\nWhat are the advantages and disadvantages of using semi-supervised learning compared to other learning paradigms?\\nHow can the quality of unlabeled data be ensured in semi-supervised learning?\\nHow can semi-supervised learning be used to address the challenge of limited labeled data in real-world applications?\\nWhat are some recent advances and state-of-the-art methods in semi-supervised learning?\\nHow can the performance of semi-supervised learning algorithms be evaluated and compared?\\nWhat are some potential ethical concerns and biases that need to be considered when using semi-supervised learning in practice?',\n",
       " 'What is reinforcement learning and how does it differ from other types of machine learning?\\nHow does reinforcement learning work, and what are the key components of a reinforcement learning system?\\nWhat are some of the most common applications of reinforcement learning in real-world scenarios?\\nHow can reinforcement learning be used to optimize decision-making processes in complex environments?\\nWhat are the challenges and limitations of reinforcement learning, and how can they be addressed?\\nHow does reinforcement learning handle the exploration-exploitation trade-off, and what are some common strategies for balancing these two objectives?\\nWhat are some of the most popular reinforcement learning algorithms, and how do they differ from one another?\\nHow can reinforcement learning be combined with other machine learning techniques, such as deep learning, to improve performance and efficiency?\\nWhat are some ethical considerations to keep in mind when deploying reinforcement learning systems in real-world applications?\\nHow is reinforcement learning being used in cutting-edge research and development, and what are some exciting new developments in this field?',\n",
       " 'What is deep learning and how does it differ from traditional machine learning?\\nWhat are the key components of a deep learning model?\\nCan you explain the concept of backpropagation in deep learning?\\nHow do convolutional neural networks (CNNs) work in image recognition tasks?\\nWhat is the role of activation functions in deep learning?\\nHow can we prevent overfitting in deep learning models?\\nCan you explain the concept of transfer learning in deep learning?\\nWhat are some popular deep learning frameworks and libraries?\\nHow is deep learning being used in natural language processing (NLP)?\\nWhat are some ethical considerations when implementing deep learning models?',\n",
       " 'What are the key components of a neural network?\\nHow do neural networks learn and adapt to new data?\\nWhat are the different types of neural networks and their applications?\\nHow can neural networks be used for image and speech recognition?\\nWhat are the challenges in designing and training deep neural networks?\\nHow can neural networks be optimized for better performance and accuracy?\\nWhat are the ethical considerations when using neural networks for decision-making?\\nHow can neural networks be used for natural language processing and generation?\\nWhat are the latest advancements and trends in neural network research?\\nHow can neural networks be integrated with other machine learning techniques for improved results?',\n",
       " 'What is natural language processing and how does it work?\\nWhat are some common applications of natural language processing in everyday life?\\nHow can natural language processing be used to improve customer service?\\nWhat are the key challenges in developing accurate natural language processing systems?\\nHow does natural language processing differ from other forms of artificial intelligence?\\nWhat are some of the most popular natural language processing techniques and algorithms?\\nHow can natural language processing be used to analyze social media data?\\nWhat are the ethical considerations when using natural language processing to analyze personal data?\\nHow can natural language processing be used to improve language translation services?\\nWhat is the future of natural language processing and how will it continue to evolve?',\n",
       " 'What is computer vision and how does it differ from image processing?\\nWhat are the key components of a computer vision system?\\nHow is machine learning used in computer vision?\\nWhat are some common applications of computer vision in industry?\\nHow does computer vision enable autonomous vehicles to navigate?\\nWhat are the challenges in developing accurate and reliable computer vision systems?\\nHow is computer vision used in medical imaging and diagnostics?\\nWhat are some popular datasets and benchmarks used in computer vision research?\\nHow does computer vision intersect with other fields such as natural language processing and robotics?\\nWhat are the ethical considerations in deploying computer vision systems, particularly in areas such as surveillance and facial recognition?',\n",
       " 'What are the different types of recommendation systems?\\nHow do recommendation systems use collaborative filtering?\\nHow do recommendation systems use content-based filtering?\\nWhat are the benefits of using a hybrid recommendation system?\\nHow can you evaluate the performance of a recommendation system?\\nWhat are some common challenges in building recommendation systems?\\nHow do recommendation systems handle the cold start problem?\\nHow can recommendation systems incorporate user feedback?\\nWhat are some ethical considerations when building recommendation systems?\\nHow do recommendation systems impact user behavior and decision-making?',\n",
       " 'What is anomaly detection and how does it work?\\nWhat are the different types of anomalies that can be detected?\\nWhat are some common algorithms used for anomaly detection?\\nHow can anomaly detection be used in cybersecurity?\\nWhat are the challenges in implementing anomaly detection in real-world scenarios?\\nHow can anomaly detection be used in predictive maintenance for industrial equipment?\\nWhat is the role of machine learning in anomaly detection?\\nHow can anomaly detection be used in fraud detection for financial transactions?\\nWhat are the best practices for evaluating the performance of anomaly detection models?\\nHow can anomaly detection be used in healthcare for early detection of diseases?']"
      ]
     },
     "execution_count": 184,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "question_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The list for each question is now collected into a single long list. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "metadata": {},
   "outputs": [],
   "source": [
    "question_list_formatted = []\n",
    "\n",
    "for question_set in question_list:\n",
    "    question_list_formatted += question_set.split(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 186,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(question_list_formatted)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating Responses from Question List\n",
    "\n",
    "Using the question list, Nemotron-4 340B Instruct can be used to generate responses to the questions. \n",
    "\n",
    "The first things needed is a function that will be used to generate the response from [build.nvidia.com](https://build.nvidia.com/)!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def generate_responses(client, question):\n",
    "    prompt = RESPONSE_PROMPT_TEMPLATE.format(question=question)\n",
    "    response = await client.chat.completions.create(\n",
    "        model=\"nvidia/nemotron-4-340b-instruct\",\n",
    "        messages=[\n",
    "            {\"role\" : \"user\",\n",
    "             \"content\" : prompt}\n",
    "        ],\n",
    "        temperature=0.2,\n",
    "        top_p=0.7,\n",
    "        max_tokens=1024,\n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, the `asycio` library allows efficient use of the API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 217,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def response_generator(client, question_list):\n",
    "    tasks = [generate_responses(client, question) for question in question_list]\n",
    "    response_list = await asyncio.gather(*tasks)\n",
    "    return response_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 218,
   "metadata": {},
   "outputs": [],
   "source": [
    "question_response_list = asyncio.run(response_generator(client, question_list_formatted))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to move to the next stage, a dataset will be created in `.jsonl` format and will store questions with the responses generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 219,
   "metadata": {},
   "outputs": [],
   "source": [
    "question_response_pair_list = []\n",
    "for question, response_set in zip(question_list_formatted, question_response_list):\n",
    "    question_response_pair_list.append(\n",
    "        {\n",
    "            \"question\" : question, \n",
    "            \"responses\" : {\n",
    "                \"response_a\" : {\"response\" : response_set.split(\"RESPONSE B:\")[0].replace(\"RESPONSE A:\", \"\").strip()},\n",
    "                \"response_b\" : {\"response\" : response_set.split(\"RESPONSE B:\")[-1].split(\"\\n\\n\")[0].strip()}\n",
    "            },\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset will be written out to a file called `synthetic_data.jsonl` below!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open('synthetic_data.jsonl', 'w') as f:\n",
    "    for item in question_response_pair_list:\n",
    "        f.write(json.dumps(item))\n",
    "        f.write('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Nemotron-4 340B Reward to Generate a Preference Dataset\n",
    "\n",
    "Equipped with a dataset that has questions that have response pairs, a preference dataset that is compatible with DPO training, SteerLM reward model training, and RLHF reward model training can be generated straightforwardly thanks to [Nemotron-4 340B Reward](https://build.nvidia.com/nvidia/nemotron-4-340b-reward) available through [build.nvidia.com](https://build.nvidia.com/)!\n",
    "\n",
    "First, an example of how to use the endpoint.\n",
    "\n",
    "1. You must both provide a user message, and an assistant message!\n",
    "2. It will return a chat-style message with the scores, as well as the scores in the `logprogs` parameter.\n",
    "\n",
    "The response package will include scores related to five attributes:\n",
    "\n",
    "1. Helpfulness: Overall helpfulness of the response to the prompt.\n",
    "2. Correctness: Inclusion of all pertinent facts without errors.\n",
    "3. Coherence: Consistency and clarity of expression.\n",
    "4. Complexity: Intellectual depth required to write response (i.e. whether the response can be written by anyone with basic language competency or requires deep domain expertise).\n",
    "5. Verbosity: Amount of detail included in the response, relative to what is asked for in the prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 221,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\n",
    "        \"role\" : \"user\",\n",
    "        \"content\" : \"Hello!\"\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"assistant\",\n",
    "        \"content\": \"Hello! How can I help you today?\"\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 222,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await client.chat.completions.create(\n",
    "        model=\"nvidia/nemotron-4-340b-reward\",\n",
    "        messages=messages,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 223,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatCompletion(id='07548838-8ef6-4feb-a1ee-66dd97905b72', choices=[Choice(finish_reason='length', index=0, logprobs=ChoiceLogprobs(content=[ChatCompletionTokenLogprob(token='helpfulness', bytes=None, logprob=4.09375, top_logprobs=[]), ChatCompletionTokenLogprob(token='correctness', bytes=None, logprob=4.03125, top_logprobs=[]), ChatCompletionTokenLogprob(token='coherence', bytes=None, logprob=4.25, top_logprobs=[]), ChatCompletionTokenLogprob(token='complexity', bytes=None, logprob=0.5703125, top_logprobs=[]), ChatCompletionTokenLogprob(token='verbosity', bytes=None, logprob=1.109375, top_logprobs=[])]), message=[ChatCompletionMessage(content='helpfulness:4.09375,correctness:4.03125,coherence:4.25,complexity:0.5703125,verbosity:1.109375', role='assistant', function_call=None, tool_calls=None)])], created=None, model=None, object=None, service_tier=None, system_fingerprint=None, usage=CompletionUsage(completion_tokens=1, prompt_tokens=54, total_tokens=55))"
      ]
     },
     "execution_count": 223,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `logprobs` can be handled in a similar fashion to message content, as demonstrated below!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 224,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ChatCompletionTokenLogprob(token='helpfulness', bytes=None, logprob=4.09375, top_logprobs=[]),\n",
       " ChatCompletionTokenLogprob(token='correctness', bytes=None, logprob=4.03125, top_logprobs=[]),\n",
       " ChatCompletionTokenLogprob(token='coherence', bytes=None, logprob=4.25, top_logprobs=[]),\n",
       " ChatCompletionTokenLogprob(token='complexity', bytes=None, logprob=0.5703125, top_logprobs=[]),\n",
       " ChatCompletionTokenLogprob(token='verbosity', bytes=None, logprob=1.109375, top_logprobs=[])]"
      ]
     },
     "execution_count": 224,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.choices[0].logprobs.content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's useful to define a simple helper function that can extract the scores to be used in the construction of a dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 225,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_scores_from_response(openai_response_template):\n",
    "    logprobs = openai_response_template.choices[0].logprobs.content\n",
    "    score_dict = {}\n",
    "    for score in logprobs:\n",
    "        score_dict[score.token] = score.logprob\n",
    "    return score_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 226,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'helpfulness': 4.09375,\n",
       " 'correctness': 4.03125,\n",
       " 'coherence': 4.25,\n",
       " 'complexity': 0.5703125,\n",
       " 'verbosity': 1.109375}"
      ]
     },
     "execution_count": 226,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_scores_from_response(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similar to the synthetic data generation above, using `asyncio` will help provide scores in a time-efficient manner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 227,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def get_response_and_scores(client, model, question, response_content):\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": question\n",
    "        },\n",
    "        {\n",
    "            \"role\": \"assistant\",\n",
    "            \"content\": response_content\n",
    "        },\n",
    "    ]\n",
    "    \n",
    "    response = await client.chat.completions.create(\n",
    "        model=model,\n",
    "        messages=messages,\n",
    "    )\n",
    "    \n",
    "    scores = get_scores_from_response(response)\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copying the list is important to avoid overwriting or modifying the original data - though it can be reloaded from `JSONL`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 228,
   "metadata": {},
   "outputs": [],
   "source": [
    "question_response_score_list = question_response_pair_list.copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Scores are calculated efficiently using `asyncio`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 231,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def process_question_response_pairs(client, model, question_response_score_list):\n",
    "    tasks = []\n",
    "    for question_response_pair in question_response_score_list:\n",
    "        question = question_response_pair[\"question\"]\n",
    "        \n",
    "        task_a = get_response_and_scores(client, model, question, question_response_pair[\"responses\"][\"response_a\"][\"response\"])\n",
    "        task_b = get_response_and_scores(client, model, question, question_response_pair[\"responses\"][\"response_b\"][\"response\"])\n",
    "        \n",
    "        tasks.append((task_a, question_response_pair, \"response_a\"))\n",
    "        tasks.append((task_b, question_response_pair, \"response_b\"))\n",
    "    \n",
    "    results = await asyncio.gather(*[task[0] for task in tasks])\n",
    "    \n",
    "    for i, (result, task_info) in enumerate(zip(results, tasks)):\n",
    "        _, question_response_pair, response_key = task_info\n",
    "        question_response_pair[\"responses\"][response_key].update(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nothing left to do but fire it off!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 232,
   "metadata": {},
   "outputs": [],
   "source": [
    "await process_question_response_pairs(client, \"nvidia/nemotron-4-340b-reward\", question_response_score_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Quality can be relatively preserved by only keeping rows that have at least a `3.0` in the overall metric - in this case helpfulness. This will help ensure that the data remains high quality. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 239,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold = 3.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "FInally, the dataset can be exported in `.JSONL` format for use in [NeMo Aligner](https://github.com/NVIDIA/NeMo-Aligner)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 240,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'synthetic_data_with_scores_filtered-{threshold}.jsonl', 'w') as f:\n",
    "    for item in question_response_score_list:\n",
    "        question = item[\"question\"]\n",
    "        response_a = item[\"responses\"][\"response_a\"]\n",
    "        response_b = item[\"responses\"][\"response_b\"]\n",
    "        response_a[\"question\"] = question\n",
    "        response_b[\"question\"] = question\n",
    "        if response_a[\"helpfulness\"] < threshold and response_b[\"helpfulness\"] < threshold:\n",
    "            continue\n",
    "        f.write(json.dumps(response_a))\n",
    "        f.write('\\n')\n",
    "        f.write(json.dumps(response_b))\n",
    "        f.write('\\n')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nvidia-sdg",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}