Thinking_Augmented_Generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Thinking Augmented Generation\n",
    "Author: [Zain Hasan](https://x.com/ZainHasan6)\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/togethercomputer/together-cookbook/blob/main/Agents/Thinking_Augmented_Generation.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "In this notebook we will explore how you can improve the quality of smaller specialized models by using reasoning models.\n",
    "\n",
    "Specifically we will get `DeepSeek-R1` to reason about a prompt and then provide the `thinking` tokens to a smaller model like `Mistral Small 3` to generate a better response.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from together import Together\n",
    "\n",
    "client = Together(api_key = \"---\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Small model alone:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "To determine the number of \"r's\" in the words \"strawberry\" and \"burberry\" combined, we need to count the \"r's\" in each word and then add them together.\n",
      "\n",
      "- The word \"strawberry\" has 3 \"r's\".\n",
      "- The word \"burberry\" has 2 \"r's\".\n",
      "\n",
      "Adding these together:\n",
      "\n",
      "3 (from strawberry) + 2 (from burberry) = 5\n",
      "\n",
      "So, there are 5 \"r's\" in the words \"strawberry\" and \"burberry\" combined.\n"
     ]
    }
   ],
   "source": [
    "question = \"How many r's are in the word strawberry and burberry combined?\"\n",
    "\n",
    "answer = client.chat.completions.create(\n",
    "  model=\"mistralai/Mistral-Small-24B-Instruct-2501\",\n",
    "  messages=[{\"role\": \"user\", \n",
    "             \"content\": question}],\n",
    ")\n",
    "\n",
    "\n",
    "print(answer.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get R1 to think about the question and then provide the thinking tokens to a smaller model like `Mistral Small 3` to generate a better response.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<think>\n",
      "Okay, let's see. The question is asking how many times the letter \"r\" appears in the words \"strawberry\" and \"burberry\" combined. Hmm, I need to make sure I count each \"r\" in both words correctly. Let me start by writing down each word separately and then check each letter one by one.\n",
      "\n",
      "First, let's take \"strawberry\". Let me spell it out: S-T-R-A-W-B-E-R-R-Y. Wait, let me count the letters. S (1), T (2), R (3), A (4), W (5), B (6), E (7), R (8), R (9), Y (10). So \"strawberry\" has 10 letters. Now, looking for the letter \"r\". Let's go through each letter again:\n",
      "\n",
      "1. S - no\n",
      "2. T - no\n",
      "3. R - yes, that's the first \"r\"\n",
      "4. A - no\n",
      "5. W - no\n",
      "6. B - no\n",
      "7. E - no\n",
      "8. R - second \"r\"\n",
      "9. R - third \"r\"\n",
      "10. Y - no\n",
      "\n",
      "Wait, hold on. So in \"strawberry\", there are three \"r\"s? Let me confirm. The spelling is S-T-R-A-W-B-E-R-R-Y. So after the B, it's E, then R, R, Y. So positions 8 and 9 are both R's. So that's two R's after the E. But the third letter is also R. So that's three R's total in \"strawberry\"? Wait, no. Let me check again. The word is \"strawberry\". Let me break it down: \"straw\" + \"berry\". \"Straw\" is S-T-R-A-W, which has one R. Then \"berry\" is B-E-R-R-Y, which has two R's. So combined, that's 1 + 2 = 3 R's in \"strawberry\". Okay, that makes sense.\n",
      "\n",
      "Now, moving on to \"burberry\". Let's spell that out: B-U-R-B-E-R-R-Y. Let me count the letters. B (1), U (2), R (3), B (4), E (5), R (6), R (7), Y (8). So \"burberry\" has 8 letters. Now, let's check for \"r\"s:\n",
      "\n",
      "1. B - no\n",
      "2. U - no\n",
      "3. R - first \"r\"\n",
      "4. B - no\n",
      "5. E - no\n",
      "6. R - second \"r\"\n",
      "7. R - third \"r\"\n",
      "8. Y - no\n",
      "\n",
      "So in \"burberry\", there are three \"r\"s as well. Let me confirm by breaking it down. The brand name \"Burberry\" is often spelled B-U-R-B-E-R-R-Y. So after the B-U-R, then another B-E-R-R-Y. So the R's are at positions 3, 6, and 7. That's three R's. So \"burberry\" has three R's.\n",
      "\n",
      "Now, adding both words together: \"strawberry\" has 3 R's and \"burberry\" has 3 R's. So 3 + 3 = 6. Therefore, the total number of R's in both words combined is 6.\n",
      "\n",
      "Wait a second, let me double-check to make sure I didn't miscount. For \"strawberry\": S-T-R-A-W-B-E-R-R-Y. The R's are at positions 3, 8, and 9. That's three. For \"burberry\": B-U-R-B-E-R-R-Y. R's at positions 3, 6, 7. That's three. So 3 + 3 = 6. Yeah, that seems right. I don't think I missed any. Sometimes when letters are close together, it's easy to miscount, but breaking it down step by step helps. So the answer should be 6.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "thought = client.chat.completions.create(\n",
    "  model=\"deepseek-ai/DeepSeek-R1\",\n",
    "  messages=[{\"role\": \"user\", \"content\": question}],\n",
    "  stop = ['</think>'] # Stop generation when </think> is encountered\n",
    ")\n",
    "\n",
    "print(thought.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prompt template to pass in the thinking tokens to the smaller model:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROMPT_TEMPLATE = \"\"\"\n",
    "Question: {question}\n",
    "Thought process: {thinking_tokens} </think>\n",
    "Answer:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The total number of 'r's in the words \"strawberry\" and \"burberry\" combined is 6.\n"
     ]
    }
   ],
   "source": [
    "answer = client.chat.completions.create(\n",
    "  model=\"mistralai/Mistral-Small-24B-Instruct-2501\",\n",
    "  messages=[{\"role\": \"user\", \n",
    "             \"content\": PROMPT_TEMPLATE.format(question = question,\n",
    "                                               thinking_tokens=thought.choices[0].message.content\n",
    "                                               )}],\n",
    ")\n",
    "\n",
    "print(answer.choices[0].message.content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}