Adding Generic Reasoning Budget (#23)

- Adding Generic Reasoning Budget (c0f14ee90b57a8218f7c6170725c2ef6e5489944)


Co-authored-by: Chris Alexiuk <[email protected]>

README.md

@@ -388,6 +388,117 @@ python3 -m sglang.launch_server --model-path nvidia/NVIDIA-Nemotron-3-Nano-30B-A
   --reasoning-parser nano_v3
 ```
 
+#### Using Budget Control
+
+The thinking budget allows developers to keep accuracy high and meet response‑time targets \- which is especially crucial for customer support, autonomous agent steps, and edge devices where every millisecond counts.
+
+With budget control, you can set a limit for internal reasoning:
+
+* `reasoning_budget`: This is a threshold that will attempt to end the reasoning trace at the next newline encountered in the reasoning trace. If no newline is encountered within 500 tokens, it will abruptly end the reasoning trace at `reasoning_budget + 500`.
+
+> NOTE: This client will work with any OpenAI API compatible endpoint.
+
+Client for supporting budget control:
+
+```py
+from typing import Any, Dict, List
+
+import openai
+from transformers import AutoTokenizer
+
+
+class ThinkingBudgetClient:
+   def __init__(self, base_url: str, api_key: str, tokenizer_name_or_path: str):
+       self.base_url = base_url
+       self.api_key = api_key
+       self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
+       self.client = openai.OpenAI(base_url=self.base_url, api_key=self.api_key)
+
+
+   def chat_completion(
+       self,
+       model: str,
+       messages: List[Dict[str, Any]],
+       reasoning_budget: int = 512,
+       max_tokens: int = 1024,
+       **kwargs,
+   ) -> Dict[str, Any]:
+       assert (
+           max_tokens > reasoning_budget
+       ), f"thinking budget must be smaller than maximum new tokens. Given {max_tokens=} and {reasoning_budget=}"
+
+
+       # 1. first call chat completion to get reasoning content
+       response = self.client.chat.completions.create(
+           model=model, messages=messages, max_tokens=reasoning_budget, **kwargs
+       )
+       content = response.choices[0].message.content
+
+
+       reasoning_content = content
+       if not "</think>" in reasoning_content:
+           # reasoning content is too long, closed with a period (.)
+           reasoning_content = f"{reasoning_content}.\n</think>\n\n"
+       reasoning_tokens_len = len(
+           self.tokenizer.encode(reasoning_content, add_special_tokens=False)
+       )
+       remaining_tokens = max_tokens - reasoning_tokens_len
+       assert (
+           remaining_tokens > 0
+       ), f"remaining tokens must be positive. Given {remaining_tokens=}. Increase the max_tokens or lower the reasoning_budget."
+
+
+       # 2. append reasoning content to messages and call completion
+       messages.append({"role": "assistant", "content": reasoning_content})
+       prompt = self.tokenizer.apply_chat_template(
+           messages,
+           tokenize=False,
+           continue_final_message=True,
+       )
+       response = self.client.completions.create(
+           model=model, prompt=prompt, max_tokens=remaining_tokens, **kwargs
+       )
+
+
+       response_data = {
+           "reasoning_content": reasoning_content.strip().strip("</think>").strip(),
+           "content": response.choices[0].text,
+           "finish_reason": response.choices[0].finish_reason,
+       }
+       return response_data
+```
+
+Calling the server with a budget (Restricted to 32 tokens here as an example)
+
+```py
+tokenizer_name_or_path = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
+client = ThinkingBudgetClient(
+   base_url="http://localhost:8000/v1",  # Nemotron 3 Nano deployed in thinking mode
+   api_key="EMPTY",
+   tokenizer_name_or_path=tokenizer_name_or_path,
+)
+
+
+result = client.chat_completion(
+   model="nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
+   messages=[
+       {"role": "system", "content": "You are a helpful assistant. /think"},
+       {"role": "user", "content": "What is 2+2?"},
+   ],
+   reasoning_budget=32,
+   max_tokens=512,
+   temperature=1.0,
+   top_p=1.0,
+)
+print(result)
+```
+
+You should see output similar to the following:
+
+```
+{'reasoning_content': "Okay, the user asked, What is 2+2? Let me think. Well, 2 plus 2 equals 4. That's a basic.", 'content': '2 + 2 equals **4**.\n', 'finish_reason': 'stop'}
+```
+
 ## Model Version(s)
 
 - v1.0