s1-test-time-scaling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "26a00332-1431-43f1-9d91-b21553fc08f1",
   "metadata": {},
   "source": [
    "### Experimenting budget forcing and test-time scaling from `s1: Simple test-time scaling`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6990ccc7-b349-4fc5-b247-c924e1a12e13",
   "metadata": {},
   "source": [
    "**From the paper: https://arxiv.org/pdf/2501.19393 <br/>**\n",
    "\n",
    ">Figure 3. Budget forcing with s1-32B. The model tries to stop\n",
    "after “...is 2.”, but we suppress the end-of-thinking token delimiter\n",
    "instead appending “Wait” leading s1-32B to self-correct its answer.\n",
    "\n",
    "<img src=\"https://arxiv.org/html/2501.19393v2/x4.png\" width=400 height=300>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56cf151a-1d5a-4ce3-a2ae-ce445b22dc09",
   "metadata": {},
   "source": [
    "**Interesting nuggets from s1 paper (methodology)**\n",
    "\n",
    "1. They collected a dataset of 1k examples with reasoning traces from Google Gemini model and performed SFT (supervised fine tuning).\n",
    "2. They fix response lengths by adding \"wait\" tokens in certain cases to get models to generate longer CoT's, verify, and correct itself, or halt token generation by introducing an EOT (end of thinking) token delimiter and the authors call this phenomenon as \"budget forcing.\"\n",
    "\n",
    "**Budget Forcing**\n",
    "Not to be picky or pedantic but `budget forcing (BF)` is still not a parallel inference scaling technique (as seen in o-1 or by Gemini Thinking). As the authors point out, we can think of `BF` as a sequential inference scaling technique. Despite the `<wait>` and `<think` tokens at appropriate steps, the model is still generating one token at a time, the only difference being in total number of tokens.\n",
    "\n",
    "These two charts highlight the limitations, but they are still salient contributions from the `s1` paper:\n",
    "<br/>\n",
    "<img src=\"https://arxiv.org/html/2501.19393v2/x7.png\" width=500 height=500 />\n",
    "<br/>\n",
    "<img src=\"https://arxiv.org/html/2501.19393v2/x8.png\" width=500 height=500 />\n",
    "</br>\n",
    "The important contribution here is `BF` method seems to be effective than other resource hungry inference-scaling techniques. A good test on how useful `BF` here could be is by testing a good instruct model (Claude or open weight models like Qwen 2.5) using CoT prompting and observe the quality of responses supplied by models trained with `BF` vs regular models with CoT prompts.\n",
    "\n",
    "**Compute optimal parallel inference scaling**\n",
    "Although `BF` is compared against majority voting, and they found it to work well there are better compute optimal strategies\n",
    "to perform inference scaling. For more check paper: https://arxiv.org/pdf/2408.03314v1, especially section **5.2**. Key strategies that should be compared against `BF` are:\n",
    "<br/>\n",
    "<img src=\"https://arxiv.org/html/2408.03314v1/x3.png\" width=500 height=500 />"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b8d26a0-5e53-4aa0-bcff-61656a479ede",
   "metadata": {},
   "source": [
    "### What is scaling at test time ?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94c33caa-6448-47ac-ae1d-b7453e5c43fc",
   "metadata": {},
   "source": [
    "Given a question `Q` to a model `LM` you get a response `R`. This process is pretty familiar and is called inference, but the key takeaway here is until `o1`, there was a semi-fixed resource consumption and also an assumption that the compute required to answer `Q`. After the advent of `o1` compute on inference or test-time compute (TTC) dynamically increase their reasoning time during inference leading to more time thinking about `Q`, improving accuracy of `R` but at the cost of higher compute usage."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a26e4b61-5068-4211-a1f1-8b475d20dfc1",
   "metadata": {},
   "source": [
    "### Isn't Deepseek-R1's scaling inference time ?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7afdf99-1a28-410c-908e-f690ecba8bbc",
   "metadata": {},
   "source": [
    "**Short answer**: Nope\n",
    "\n",
    "**Long answer**:\n",
    "DeepSeek stated that their main goal was to achieve strong reasoning capabilities by leveraging the principle of deep, step-by-step thinking. DeepSeek-R1 got nearly SoTA results on most reasoning datasets simply by leveraging pure RL (GRPO), good data quality mix, and an incredibly strong base model (DeepSeek-V3). So, they built a reasoning model that mimicks o-1, but the phenomenon of scaling at inference time is somewhat not disclosed in the paper. Also, page-1 from s1's paper: https://arxiv.org/pdf/2501.19393 <br/>\n",
    ">However, despite the large number of o1 replication attempts, none have openly replicated\n",
    "a clear test-time scaling behavior. Thus, we ask: what is\n",
    "the simplest approach to achieve both test-time scaling and\n",
    "strong reasoning performance?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b3dd0df-0229-4614-a650-66b518db3a5a",
   "metadata": {},
   "source": [
    "### Experimenting `s1` phenomenon on Science of Science tasks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7232a33-5d3e-453c-97bf-bd0f98710a6b",
   "metadata": {},
   "source": [
    "### 1. Import necessary libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "72e30633-02cf-46a8-96a5-72f2cccb535d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This Source Code Form is subject to the terms of the MIT\n",
    "# License. If a copy of the same was not distributed with this\n",
    "# file, You can obtain one at\n",
    "# https://github.com/akhilpandey95/s1/blob/main/LICENSE.\n",
    "\n",
    "import gc\n",
    "import time\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, set_seed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49ede556-0eed-47d7-8d63-1ddbe2ec8ad0",
   "metadata": {},
   "source": [
    "### 2. Load Deepseek-R1 distill models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adcb6616-fd66-48b9-99b5-9706a57079c5",
   "metadata": {},
   "source": [
    "#### 2.1 Model, tokenizer load"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "33180287-d83e-4be1-bc22-d6c8cf96be05",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "del model, tokenizer\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "439958a3-1168-4e6f-9ff9-b5c96ff8fc27",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------\n",
      "Using cpu to load /Users/akhilakella/code/models/DeepSeek-R1-Distill-Qwen-1.5B\n",
      "----------------------------------\n",
      "Model-tokenizer Load Time: 0.6631889343261719\n",
      "----------------------------------\n"
     ]
    }
   ],
   "source": [
    "# device for acceleration\n",
    "# device = \"mps\" if torch.mps.is_available() else \"cpu\"\n",
    "device = \"cpu\"\n",
    "\n",
    "# set the model-id\n",
    "model_id = \"/Users/akhilakella/code/models/DeepSeek-R1-Distill-Qwen-1.5B\"\n",
    "\n",
    "# log\n",
    "print(\"----------------------------------\")\n",
    "print(f\"Using {device} to load {model_id}\")\n",
    "print(\"----------------------------------\")\n",
    "\n",
    "# 4-bit quantization config\n",
    "bnb_4bit = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "    bnb_4bit_quant_storage=torch.bfloat16\n",
    ")\n",
    "\n",
    "# get model-tokenizer pair\n",
    "start = time.time()\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_id, \\\n",
    "                                             trust_remote_code=True, \\\n",
    "                                             low_cpu_mem_usage=True, \\\n",
    "                                             torch_dtype=torch.bfloat16, \\\n",
    "                                             device_map=device)\n",
    "\n",
    "# # pad token if needed\n",
    "# tokenizer.add_special_tokens({\"pad_token\": \"<|finetune_right_pad_id|>\"})\n",
    "# model.resize_token_embeddings(len(tokenizer))\n",
    "\n",
    "# load time\n",
    "end = time.time()\n",
    "print(\"Model-tokenizer Load Time:\", end - start)\n",
    "print(\"----------------------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b79f0e68-f0a4-498a-abfa-5a032a31d1fe",
   "metadata": {},
   "source": [
    "#### 2.2 Seed and parameters for reproducibility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "436c890c-2a2a-460b-8961-3f2c6cd2c2a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# seed for reproducibility\n",
    "set_seed(2025)\n",
    "\n",
    "# set top_p and temperature to none\n",
    "model.generation_config.temperature=None\n",
    "model.generation_config.top_p=None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "897cd80e-cc54-4437-984b-08c4908f4c4b",
   "metadata": {},
   "source": [
    "### 3. Prompting a reasoning model to construct a policy brief"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48d0d1d7-f6f8-4f27-ad8b-bf2589d09b47",
   "metadata": {},
   "source": [
    "#### 3.1 Prompt for T, K, and E"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e969e09e-4971-49db-961a-daae8e508d30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prompt to generate T, K, and E\n",
    "input_text = \"\"\"\n",
    "SCIENTIFIC_PAPER_TITLE: Pharmaceutical Side Effects and Mental Health Paradoxes among Racial-Ethnic Minorities\n",
    "\n",
    "SCIENTIFIC_PAPER_ABSTRACT: Sociologists have long struggled to explain the minority mental health paradox: \n",
    "that racial-ethnic minorities often report better mental health than non-Hispanic whites despite social \n",
    "environments that seem less conducive to well-being. Using data from the 2008–2013 Medical Expenditure \n",
    "Panel Survey (MEPS), this study provides a partial explanation for the paradox rooted in a very different disparity. \n",
    "Evidence from MEPS indicates that non-Hispanic whites consume more pharmaceuticals than racial-ethnic minorities \n",
    "for a wide variety of medical conditions. Moreover, non-Hispanic whites consume more pharmaceuticals that although \n",
    "effective in treating their focal indication, include depression or suicide as a side effect. In models that adjust\n",
    "for the use of such medications, the minority advantage in significant distress is reduced, in some instances to \n",
    "statistical nonsignificance. Although a significant black and Hispanic advantage in a continuous measure of distress \n",
    "remains, the magnitude of the difference is reduced considerably. The relationship between the use of medications with \n",
    "suicide as a side effect and significant distress is especially large, exceeding, for instance, the relationship between \n",
    "poverty and significant distress. For some minority groups, the less frequent use of such medications is driven by better \n",
    "health (as in the case of Asians), whereas for others, it reflects a treatment disparity (as in the case of blacks), \n",
    "although the consequences for the mental health paradox are the same. The implications of the results are discussed, \n",
    "especially with respect to the neglect of psychological side effects in the treatment of physical disease as well as \n",
    "the problem of multiple morbidities.\n",
    "\n",
    "The final policy brief must contain the following structure:\n",
    "\n",
    "1. Title\n",
    "2. Executive summary of the scientific work\n",
    "3. Key policy implications\n",
    "\n",
    "Write a policy brief for the above scientific work.\n",
    "\"\"\"\n",
    "\n",
    "# apply chat template\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": (\n",
    "            \"You are a helpful intelligent policy brief writing assistant. \"\n",
    "            \"You have the ability to take scientific papers and write policy briefs \"\n",
    "            \"for the respective scientific works.\"\n",
    "        )\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": input_text\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "88e32094-ab06-4861-9e7b-96c14e95e97b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encoded inputs: 438\n"
     ]
    }
   ],
   "source": [
    "# get attention mask and input ids\n",
    "inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, \\\n",
    "                                        return_dict=True, return_tensors=\"pt\")\n",
    "input_encoded = {k: v.to(model.device) for k, v in inputs.items()}\n",
    "input_shape = len(input_encoded[\"input_ids\"][0])\n",
    "\n",
    "# set max new tokens\n",
    "max_tokens = 1024\n",
    "\n",
    "# shape check\n",
    "print(f\"Encoded inputs: {input_shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ca4cdad-cad9-4894-9f77-e8861ae0cfb0",
   "metadata": {},
   "source": [
    "#### 3.3 Inference on the prompt using `R1-distill-Qwen-8b` without `BF`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "76e9399f-09dc-4c19-a7f9-87c1934e2232",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------\n",
      "Policy Brief:\n",
      "Okay, so I need to write a policy brief based on the provided scientific paper. Let me start by understanding the structure and the content of the paper.\n",
      "\n",
      "The paper is titled \"Pharmaceutical Side Effects and Mental Health Paradoxes Among Racial-Ethnic Minorities\" and it's about the mental health paradox where racial-ethnic minorities report better mental health than non-Hispanic whites despite having less favorable social environments. The study uses data from the 2008-2013 Medical Expenditure Panel Survey (MEPS) to explore this issue.\n",
      "\n",
      "First, I need to create a title that reflects the study's focus. Maybe something like \"Exploring the Mental Health Paradox Among Racial-Ethnic Minorities: The Role of Pharmaceutical Side Effects.\"\n",
      "\n",
      "Next, the executive summary should summarize the key findings. The study found that racial-ethnic minorities consume more pharmaceuticals, including those that cause depression or suicide, which reduces their mental health advantage. However, there's a significant link between the use of such medications and significant distress, especially among Asian minorities. The paper also discusses the implications for mental health care, particularly regarding psychological side effects and the management of multiple conditions.\n",
      "\n",
      "For the key policy implications, I should suggest actions that address these findings. For example, implementing policies that reduce the use of medications with side effects like depression or suicide could help mitigate the paradox. Also, providing better mental health resources for minority groups might improve their mental health outcomes. Additionally, encouraging the use of alternative treatments or medications that don't cause significant distress could be beneficial.\n",
      "\n",
      "I should also consider the limitations of the study. It's based on a specific survey, so the findings might not be generalizable to all populations. There might be other factors influencing the paradox that weren't measured in the study. Therefore, the policy recommendations should be cautious and based on the evidence presented.\n",
      "\n",
      "I need to make sure the policy brief is clear and concise, providing both the findings and actionable recommendations. It should also highlight the importance of addressing these issues to improve mental health outcomes for all racial-ethnic minorities.\n",
      "\n",
      "Let me structure this step by step. Start with the title, then the executive summary, followed by the key policy implications. I'll make sure each section flows logically into the next, providing a comprehensive overview of the study and its implications for policy.\n",
      "\n",
      "I should also ensure that the language is accessible, avoiding overly technical terms unless necessary. The goal is to inform policymakers about the findings and suggest practical steps they can take up.\n",
      "\n",
      "Finally, I'll review the policy brief to make sure it adheres to the required structure and that all key points are covered without being too verbose.\n",
      "</think>\n",
      "\n",
      "**Title:** Exploring the Mental Health Paradox Among Racial-Ethnic Minorities: The Role of Pharmaceutical Side Effects\n",
      "\n",
      "**Executive Summary:**\n",
      "This study examines the mental health paradox where racial-ethnic minorities report better mental health than non-Hispanic whites, despite having less favorable social environments. Using data from the 2002-2013 Medical Expenditure Panel Survey (MEPS), the research highlights that racial-ethnic minorities consume more pharmaceuticals, including those causing depression or suicide, which reduces their mental health advantage. However, there's a significant link between the use of such medications and significant distress, particularly among Asian minorities. The study also discusses the implications of these findings for mental health care, emphasizing the need to address psychological side effects and the management of multiple conditions.\n",
      "\n",
      "**Key Policy Implications:**\n",
      "\n",
      "1. **Addressing Medication Use Disparities:** Implement policies to reduce the use of medications with significant side effects like depression or suicide, which can mitigate the mental health paradox.\n",
      "\n",
      "2. **Enhance Mental Health Resources:** Provide better mental health resources and support for minority groups, potentially improving their mental health outcomes.\n",
      "\n",
      "3. **Encourage Alternative Treatments:** Suggest encouraging the use of alternative or non-significant medications to reduce the mental health paradox.\n",
      "\n",
      "4. **Promote Inclusive Mental Health Services:** Advocate for the development of inclusive mental health services tailored to the needs of minority populations.\n",
      "\n",
      "5. **Consider Generalizability:** Note that the study's findings may not be representative of all populations, suggesting the need for further research to generalize the results.\n",
      "\n",
      "**Conclusion:**\n",
      "This study underscores the importance of addressing medication use disparities and enhancing mental health resources to improve the mental health outcomes of racial-ethnic minorities. Policymakers should consider these findings when developing policies aimed at reducing the mental health paradox and promoting inclusive mental health services.<｜end▁of▁sentence｜>\n",
      "----------------------------------\n",
      "Inference time: 44.09 seconds\n",
      "New tokens generated: 919\n",
      "Tokens per second: 20.84\n"
     ]
    }
   ],
   "source": [
    "# time check\n",
    "start = time.time()\n",
    "\n",
    "# routine model.generate()\n",
    "outputs = model.generate(\n",
    "    **input_encoded,\n",
    "    max_new_tokens=max_tokens,\n",
    "    do_sample=False,\n",
    "    num_beams=1,\n",
    "    pad_token_id=tokenizer.pad_token_id,\n",
    "    eos_token_id=tokenizer.eos_token_id\n",
    ")\n",
    "\n",
    "# inference time\n",
    "end = time.time()\n",
    "inference = end - start\n",
    "\n",
    "# model.generate() output\n",
    "print(\"----------------------------------\")\n",
    "output = tokenizer.decode(outputs[0][input_shape:])\n",
    "print(\"Policy Brief:\")\n",
    "print(output)\n",
    "print(\"----------------------------------\")\n",
    "\n",
    "# tps logic\n",
    "new_tokens = outputs[0].shape[0] - input_shape\n",
    "tokens_per_second = new_tokens / inference if inference > 0 else float('inf')\n",
    "print(f\"Inference time: {inference:.2f} seconds\")\n",
    "print(f\"New tokens generated: {new_tokens}\")\n",
    "print(f\"Tokens per second: {tokens_per_second:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bb40818-5bc1-45ad-970d-d08bebf91ca2",
   "metadata": {},
   "source": [
    "#### 3.4 Inference on the prompt with test time scaling via `BF`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ddeb474d-0c1b-4322-992e-269db09e6608",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "287"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "del model, tokenizer\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "525ea223-e951-429f-93a1-880ae6b0cc85",
   "metadata": {},
   "source": [
    "##### 3.4.1 Induce **Wait** and `</think>` tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "29e5fc85-d813-488b-b65c-be1b28600bf0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------\n",
      "Before adjusting tokens: ['<｜begin▁of▁sentence｜>', '<｜end▁of▁sentence｜>']\n",
      "----------------------------------\n",
      "After adding special wait, think tokens: ['<｜begin▁of▁sentence｜>', '<｜end▁of▁sentence｜>', 'Wait', '</think>']\n",
      "----------------------------------\n",
      "wait_token: Wait\n",
      "end_think_token_id: 151649\n",
      "----------------------------------\n"
     ]
    }
   ],
   "source": [
    "# checking the tokenizer's special tokens\n",
    "print(\"----------------------------------\")\n",
    "print(f\"Before adjusting tokens: {tokenizer.all_special_tokens}\")\n",
    "print(\"----------------------------------\")\n",
    "\n",
    "# define BF parameters\n",
    "bf_tokens = 0\n",
    "wait_token = \"Wait\"\n",
    "end_think_token = \"</think>\"\n",
    "end_think_triggered = False\n",
    "NUM_IGNORE = 1\n",
    "\n",
    "# resize model embedding\n",
    "special_tokens_dict = {\"additional_special_tokens\": [wait_token, end_think_token]}\n",
    "tokenizer.add_special_tokens(special_tokens_dict)\n",
    "model.resize_token_embeddings(len(tokenizer), mean_resizing=False)\n",
    "print(f\"After adding special wait, think tokens: {tokenizer.all_special_tokens}\")\n",
    "print(\"----------------------------------\")\n",
    "\n",
    "# check\n",
    "wait_token_id = tokenizer.convert_tokens_to_ids(\"Wait\")\n",
    "end_think_token_id = tokenizer.convert_tokens_to_ids(\"</think>\")\n",
    "print(f\"wait_token: {wait_token}\")\n",
    "print(f\"end_think_token_id: {end_think_token_id}\")\n",
    "print(\"----------------------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11a8f219-0eb0-4bba-8974-73ec839c7e69",
   "metadata": {},
   "source": [
    "##### 3.4.2 Generate tokens with a budget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "891f1165-2679-4dd6-bbde-4ca61feab339",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Okay, so I need to write a policy brief based on the provided scientific paper. The title is already given, so I don't need to change that. The executive summary should summarize the study, key policy implications should discuss the findings, and the policy brief should outline the implications for policymakers.\n",
      "\n",
      "First, I'll start with the title. It's about the paradox of pharmaceutical side effects and mental health among racial-ethnic minorities. That's clear, but maybe I can make it a bit more specific, like mentioning the specific medications or the context of the study.\n",
      "\n",
      "Next, the executive summary. I need to explain the paradox: why minorities have better mental health despite lower access to medications. The study uses MEPS data, which is a big survey, so I should mention that. It found that whites consume more medications, especially those with mental health issues. Also, some medications have side effects like depression or suicide, which might explain the paradox. The study adjusts for these side effects, so the main advantage for minorities is in continuous measures like anxiety or depression, but not as much in significant distress. There's a significant link between using certain medications with side effects and significant distress, especially for blacks and Asians. So, the policy implications should suggest that mental health interventions should be tailored to the specific needs of minority groups, especially considering the side effects of medications.\n",
      "\n",
      "I should structure the policy brief with three sections: title, executive summary, key policy implications, and a conclusion. The conclusion should tie everything together, emphasizing the need for targeted interventions.\n",
      "\n",
      "Wait\n",
      "----------------------------------\n",
      "Synchronous BF Inference\n",
      "Inference time: 1741.78 seconds\n",
      "New tokens generated: 0\n",
      "----------------------------------\n",
      ",,,,"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "KeyboardInterrupt\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# clone the input ids\n",
    "generated = input_encoded[\"input_ids\"].clone()\n",
    "\n",
    "# sequential generation of tokens\n",
    "start = time.time()\n",
    "while bf_tokens < max_tokens:\n",
    "    prompt_text = tokenizer.decode(generated[0])\n",
    "    input_encoded = tokenizer(prompt_text, return_tensors=\"pt\").to(model.device)\n",
    "\n",
    "    # generate one token at a time\n",
    "    outputs = model.generate(\n",
    "        **input_encoded,\n",
    "        max_new_tokens=1,\n",
    "        do_sample=False,\n",
    "        num_beams=1,\n",
    "        pad_token_id=tokenizer.pad_token_id,\n",
    "        eos_token_id=tokenizer.eos_token_id\n",
    "    )\n",
    "\n",
    "    # \n",
    "    next_token = outputs[0][-1].unsqueeze(0).unsqueeze(0)\n",
    "    generated = torch.cat((generated, next_token), dim=1)\n",
    "    token_str = tokenizer.decode(next_token[0])\n",
    "\n",
    "    # display the text generated\n",
    "    print(token_str, end=\"\", flush=True)\n",
    "\n",
    "    # increment token counter\n",
    "    bf_tokens += 1\n",
    "\n",
    "    # check for wait or think tokens\n",
    "    if wait_token in token_str:\n",
    "        # out of the loop\n",
    "        break\n",
    "    if end_think_token in token_str:\n",
    "        # end of token triggered\n",
    "        end_think_triggered = True\n",
    "\n",
    "        # out of the loop\n",
    "        break\n",
    "\n",
    "end = time.time()\n",
    "inference_time = end - start\n",
    "total_generated = generated.shape[1] - len(input_encoded[\"input_ids\"][0])\n",
    "\n",
    "print(\"\\n----------------------------------\")\n",
    "print(\"Synchronous BF Inference\")\n",
    "print(f\"Inference time: {inference_time:.2f} seconds\")\n",
    "print(f\"New tokens generated: {total_generated}\")\n",
    "print(\"----------------------------------\")\n",
    "\n",
    "# If a wait token was encountered, ignore the stop token NUM_IGNORE times\n",
    "if wait_token in tokenizer.decode(generated[0]):\n",
    "    for i in range(NUM_IGNORE):\n",
    "        # Adjust remaining token budget (here simply using bf_tokens)\n",
    "        remaining_tokens = max_tokens - bf_tokens\n",
    "        # Append generated text and the ignore string to the prompt\n",
    "        prompt = tokenizer.decode(generated[0]) + wait_token\n",
    "        input_encoded = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
    "        # Generate additional tokens synchronously until a control token is seen\n",
    "        while bf_tokens < max_tokens:\n",
    "            outputs = model.generate(\n",
    "                **input_encoded,\n",
    "                max_new_tokens=1,\n",
    "                do_sample=False,\n",
    "                num_beams=1,\n",
    "                pad_token_id=tokenizer.pad_token_id,\n",
    "                eos_token_id=tokenizer.eos_token_id\n",
    "            )\n",
    "            next_token = outputs[0][-1].unsqueeze(0).unsqueeze(0)\n",
    "            generated = torch.cat((generated, next_token), dim=1)\n",
    "            token_str = tokenizer.decode(next_token[0])\n",
    "            print(token_str, end=\"\", flush=True)\n",
    "            bf_tokens += 1\n",
    "            if wait_token in token_str or end_think_token in token_str:\n",
    "                break\n",
    "\n",
    "# force end_think_token\n",
    "if not end_think_triggered:\n",
    "    prompt = tokenizer.decode(generated[0]) + end_think_token + \"\\n\"\n",
    "    input_encoded = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
    "    outputs = model.generate(\n",
    "         **input_encoded,\n",
    "         max_new_tokens=1024,\n",
    "         do_sample=False,\n",
    "         num_beams=1,\n",
    "         pad_token_id=tokenizer.pad_token_id,\n",
    "         eos_token_id=tokenizer.eos_token_id\n",
    "    )\n",
    "    final_text = tokenizer.decode(outputs[0][input_encoded[\"input_ids\"].shape[1]:])\n",
    "    print(\"\\n----------------------------------\")\n",
    "    print(\"Final Answer:\")\n",
    "    print(final_text)\n",
    "    print(\"----------------------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01c1fbdb-11d5-4ba1-81f5-8f90d612c7f6",
   "metadata": {},
   "source": [
    "### 4. Prompting a reasoning model to extract information from technology disclosure page"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d40bb0a-f3e7-435c-8027-0444e9117e69",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd9d4ade-0498-4d96-bae2-c96cea967ce1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "469dbcf6-a89c-4ba1-955c-28f62c6541a2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}