Token Budget Calculator
Calculate token budgets for a prompt with system, examples, and query.
Given a maximum context length and the token counts for each part, determine how many few-shot examples can fit while reserving space for the query and a completion buffer.
Input:
- Line 1: max_tokens completion_buffer
- Line 2: system_tokens query_tokens
- Line 3: N (number of available examples)
- Next N lines: example_tokens (tokens per example)
Output:
- Line 1: Number of examples that fit
- Line 2: Total tokens used (system + selected examples + query)
- Line 3: Tokens remaining for completion
Example
Input
4096 512 200 100 5 150 150 150 150 150
Output
3 650 2934
Reasoning
- The maximum available tokens for examples and the query are calculated by subtracting the completion buffer from the max tokens: .
- We reserve space for the system and query tokens: tokens.
- The remaining tokens for examples are , and each example takes 150 tokens, so we can fit examples, but since we have 5 examples and the calculator can only use whole examples, we calculate how many examples fit: , so 21 examples would take tokens, exceeding the available tokens when adding system and query tokens, thus we try with fewer examples.
- We try with 3 examples: tokens for examples, total tokens used, leaving tokens for completion, which fits within the completion buffer, so 3 examples fit.
Constraints
- Select examples in order (first N that fit)
- Must reserve completion_buffer tokens
- Total used = system + examples + query
- Available = max_tokens - completion_buffer