Thinking 예산

Thinking 예산 (Thinking Budget)

이 예시는 Qwen3 시리즈 모델로 thinking 예산을 사용하는 추론 과정을 보여드려요. 이 과정은 두 단계로 이루어져요: (1) 모델이 지정된 thinking 예산 안에서 추론 내용을 생성하고, (2) 추론 내용을 대화 컨텍스트에 추가한 뒤 모델을 다시 호출해 최종 응답을 얻는 방식이에요.

출처: 문서

본문

환경 설정

  • transformers >= 4.51.0
  • openai >= 1.65.0

기본 사용법

먼저 Qwen3 모델을 thinking 모드로 시작해야 해요. 자세한 내용은 퀵스타트를 참고하세요.

그 다음 아래 코드로 thinking 예산을 사용해 모델을 호출할 수 있어요.

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]],
        thinking_budget: int = 512,
        max_tokens: int = 1024,
        **kwargs
    ) -> Dict[str, Any]:
        assert max_tokens > thinking_budget, f"thinking budget must be smaller than maximum new tokens. Given {max_tokens=} and {thinking_budget=}"

        # 1. first call chat completion to get reasoning content
        response = self.client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=thinking_budget,
            **kwargs
        )

        content = response.choices[0].message.content
        reasoning_content = response.choices[0].message.reasoning_content.strip("\n")
        if content is None:
            # reasoning content is too long
            reasoning_content = (
                f"{reasoning_content}"
                "\n\nConsidering the limited time by the user, "
                "I have to give the solution based on the thinking directly now."
            )
        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 thinking_budget."

        # 2. append reasoning content to messages and call completion
        messages.append({"role": "assistant", "content": f" thinking\n{reasoning_content}\n response\n\n"})
        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,
            "content": response.choices[0].text,
            "finish_reason": response.choices[0].finish_reason,
        }
        return response_data

tokenizer_name_or_path = "Qwen/Qwen3-8B"
client = ThinkingBudgetClient(
    base_url="http://localhost:30000/v1", # Qwen3 deployed in thinking mode
    api_key="EMPTY",
    tokenizer_name_or_path=tokenizer_name_or_path
)

result = client.chat_completion(
    model="Qwen3-8B",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Tell me a funny story about a cat."}
    ],
    thinking_budget=512,
    max_tokens=1024,
)
print(result["content"])

코드에서 확인할 수 있듯이, thinking_budget은 첫 번째 호출에서 모델이 쓰는 token 수를 제한하고, max_tokens는 전체 생성 예산을 뜻해요. 첫 번째 호출에 max_tokens=thinking_budget을 전달해 추론 내용만 뽑아낸 뒤, 추론 내용을 assistant 메시지로 붙여 다시 호출해서 최종 답변을 얻는 구조예요.

더 알아보기 (Learn more)