Budget Forcing and Wait Tokens

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from typing import List, Optional
import re

# Load model
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
print(f"Loading {model_name}...")

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, dtype="auto")

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"
model = model.to(device)
model.eval()

print(f"Loaded on {device}")

def generate_with_budget_forcing(prompt: str, 
                                  min_tokens: int = 50,
                                  max_tokens: int = 200,
                                  wait_token: str = " Wait,",
                                  end_markers: List[str] = None) -> str:
    """
    Generate with budget forcing.
    
    If the model tries to stop before min_tokens, append wait_token
    to encourage continued reasoning.
    
    Args:
        prompt: The input prompt
        min_tokens: Minimum tokens to generate (budget)
        max_tokens: Maximum tokens to generate
        wait_token: Token to append for continuation
        end_markers: Phrases that indicate final answer
    
    Returns:
        Generated text with reasoning
    """
    if end_markers is None:
        end_markers = ["the answer is", "therefore", "final answer", "in conclusion"]
    
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    generated_text = ""
    total_tokens = 0
    n_waits = 0
    
    while total_tokens < max_tokens:
        # Generate a chunk
        chunk_size = min(30, max_tokens - total_tokens)
        
        with torch.no_grad():
            current_input = tokenizer(prompt + generated_text, return_tensors="pt").to(device)
            outputs = model.generate(
                **current_input,
                max_new_tokens=chunk_size,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
            )
        
        new_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        new_text = new_text[len(prompt) + len(generated_text):]
        generated_text += new_text
        total_tokens = len(tokenizer.encode(generated_text))
        
        # Check if model is trying to conclude
        lower_text = generated_text.lower()
        is_concluding = any(marker in lower_text for marker in end_markers)
        
        # If concluding too early, append wait token
        if is_concluding and total_tokens < min_tokens:
            generated_text += wait_token
            n_waits += 1
            continue
        
        # If naturally concluded and past minimum, stop
        if is_concluding:
            break
        
        # If we hit EOS, check if we need to continue
        if outputs[0][-1] == tokenizer.eos_token_id:
            if total_tokens < min_tokens:
                generated_text += wait_token
                n_waits += 1
            else:
                break
    
    return generated_text, n_waits


# Test without budget forcing
prompt = "Question: What is 15 + 28?\nAnswer: Let me solve this."

print("Without budget forcing:")
print("="*60)
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=50, temperature=0.7, do_sample=True, pad_token_id=tokenizer.eos_token_id)
short_response = tokenizer.decode(outputs[0], skip_special_tokens=True)[len(prompt):]
print(f"Response: {short_response}")
print(f"Tokens: {len(tokenizer.encode(short_response))}")

# Test WITH budget forcing
print("\nWith budget forcing (min 50 tokens):")
print("="*60)

long_response, n_waits = generate_with_budget_forcing(
    prompt, 
    min_tokens=50,
    max_tokens=150
)

print(f"Response: {long_response}")
print(f"\nTokens: {len(tokenizer.encode(long_response))}")
print(f"Wait tokens inserted: {n_waits}")

# Compare different continuation tokens
continuation_tokens = [
    " Wait,",
    " Hmm,",
    " Let me think more.",
    " Actually,"
]

base_text = "Question: What is 7 × 8?\nAnswer: Let me calculate. 7 × 8 = 54. The answer is 54."

print("How different tokens affect continuation:")
print("="*60)

for token in continuation_tokens:
    prompt = base_text + token
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=40,
            temperature=0.5,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    
    continuation = tokenizer.decode(outputs[0], skip_special_tokens=True)[len(prompt):]
    print(f"\n{token.strip()}")
    print(f"  → {continuation[:80]}...")

def estimate_difficulty(problem: str) -> float:
    """
    Simple heuristic to estimate problem difficulty.
    
    In practice, you might use a classifier or the model's
    own uncertainty estimates.
    """
    difficulty = 0.3  # Base difficulty
    
    # More numbers = harder
    numbers = re.findall(r'\d+', problem)
    difficulty += 0.1 * min(len(numbers), 5)
    
    # Larger numbers = harder
    if numbers:
        max_num = max(int(n) for n in numbers)
        if max_num > 100:
            difficulty += 0.2
        if max_num > 1000:
            difficulty += 0.2
    
    # Multiple operations = harder
    ops = len(re.findall(r'[+\-×÷*/]', problem))
    difficulty += 0.1 * min(ops, 3)
    
    # Keywords suggesting complexity
    if any(w in problem.lower() for w in ['percent', 'ratio', 'average', 'probability']):
        difficulty += 0.2
    
    return min(1.0, difficulty)


def adaptive_budget(problem: str, base_tokens: int = 30, 
                    max_tokens: int = 150) -> int:
    """
    Set token budget based on problem difficulty.
    """
    difficulty = estimate_difficulty(problem)
    budget = int(base_tokens + difficulty * (max_tokens - base_tokens))
    return budget


# Test on problems of varying difficulty
problems = [
    "What is 5 + 3?",
    "What is 15% of 80?",
    "If a train travels 120 miles in 2 hours, what is its average speed?",
    "A store has 1250 items. They sell 15% and receive 340 more. How many items do they have?",
]

print("Adaptive budget based on difficulty:")
print("="*60)

for problem in problems:
    diff = estimate_difficulty(problem)
    budget = adaptive_budget(problem)
    print(f"\nProblem: {problem[:50]}...")
    print(f"  Difficulty: {diff:.2f}")
    print(f"  Token budget: {budget}")

# Example of what a good training example looks like
example_trace = """
Question: A store has 45 items. They sell 30% of them and receive 20 more.
How many items do they have now?

Let me solve this step by step.

Step 1: Calculate 30% of 45.
30% = 0.30
0.30 × 45 = 13.5

Wait, can you sell half an item? Let me reconsider.
Actually, 30% of 45 = 0.3 × 45 = 13.5
Since we can't sell half an item, let's round to 13 or 14.

Hmm, the problem might expect exact math. Let me continue with 13.5 
and see if we need to round at the end.

Step 2: Subtract items sold.
45 - 13.5 = 31.5 items remaining

Step 3: Add new items.
31.5 + 20 = 51.5 items

Rounding to a whole number: 51 or 52 items.

The answer is approximately 51-52 items (or exactly 51.5 if fractional items are allowed).
"""

print("Example reasoning trace for training:")
print("="*60)
print(example_trace)
print("\nKey features:")
print("- Step-by-step structure")
print("- Self-correction (Wait...)")
print("- Explicit uncertainty handling")
print("- Clear final answer")