diff --git a/config/experiments/deepconf/deepconf_api_test.yaml b/config/experiments/deepconf/deepconf_api_test.yaml deleted file mode 100644 index da78ea9e..00000000 --- a/config/experiments/deepconf/deepconf_api_test.yaml +++ /dev/null @@ -1,38 +0,0 @@ -# @package _global_ -# Run: CUDA_VISIBLE_DEVICES=0 python scripts/run_tts_eval.py --config-path=../config --config-name=experiments/deepconf/deepconf_api_test -# Note: Do NOT specify CUDA_VISIBLE_DEVICES when submitting via scripts/slurm/submit.sh - SLURM assigns GPUs automatically - -# Run naming -run_name: "seed${system.seed}_${now:%H-%M-%S}" - -defaults: - - dataset: small_gsm8k - - model: together_ai - - generation: default - - output: default - - system: default - - strategy: deep_think_confidence - - _self_ - -# Override some settings for DeepConf API testing -verbose: true - -# DeepConf doesn't need a scorer like ReasonEval -# It uses internal confidence-based filtering -scorer: null - -# API-specific overrides for testing -strategy: - num_paths: 8 # Moderate number for API testing - filtering_percentage: 0.6 # Keep top 60% (more selective) - temperature: 0.8 # Higher temperature for diversity - confidence_metric: "avg_confidence" - mode: "offline" - -generation: - temperature: 0.8 # Should match strategy temperature - max_new_tokens: 300 # Shorter for faster API calls - -# Use smaller dataset for testing -dataset: - subset: 20 # Only 20 questions for testing diff --git a/config/experiments/deepconf/deepconf_vllm_qwen3_aime2025.yaml b/config/experiments/deepconf/deepconf_vllm_qwen3_aime2025.yaml deleted file mode 100644 index 98ddab64..00000000 --- a/config/experiments/deepconf/deepconf_vllm_qwen3_aime2025.yaml +++ /dev/null @@ -1,72 +0,0 @@ -# @package _global_ -# Run: CUDA_VISIBLE_DEVICES=0 python scripts/run_tts_eval.py --config-path=../config --config-name=experiments/deepconf/deepconf_vllm_qwen3_aime2025 -# Note: Do NOT specify CUDA_VISIBLE_DEVICES when submitting via scripts/slurm/submit.sh - SLURM assigns GPUs automatically -# DeepConf Offline Mode: Qwen3-8B via vLLM on AIME 2025 -# Fast inference with PagedAttention - no OOM on long sequences - -defaults: - - /config - - /dataset/aime_2025 - - /model/vllm_qwen3 - - /generation/default - - /system/default - - /evaluation/default - - _self_ - -# Run naming -run_name: "deepconf_vllm_qwen3_aime2025_seed${system.seed}_${now:%H-%M-%S}" - -# Main configuration -verbose: true -report_to: null # Disable wandb for local runs - -# Model configuration (vLLM backend) -model: - type: "vllm" - model_path: "Qwen/Qwen3-8B" - model_short_name: "qwen3_8b" - gpu_memory_utilization: 0.9 - tensor_parallel_size: 1 - enable_prefix_caching: true - trust_remote_code: true - max_context_budget: 32768 - disable_thinking_mode: true # Disable Qwen3 thinking mode - -# System configuration -system: - device: cuda - seed: 42 - -# Generation configuration -generation: - max_new_tokens: 32768 # AIME problems require more tokens - temperature: 0.6 - top_p: 0.95 - batch_size: 1 - max_length: 8192 - -# No scorer needed (DeepConf has built-in weighted voting) -scorer: null - -# DeepConf offline mode strategy -strategy: - type: deepconf - mode: "offline" # Generate all traces, filter by confidence, weighted vote - budget: 16 # Number of reasoning traces to generate (all in ONE vLLM call!) - - # Generation parameters - window_size: 2048 # Sliding window for confidence computation - temperature: 0.6 # Sampling temperature for diversity - top_p: 0.95 # Nucleus sampling - max_tokens: 32768 # vLLM handles long sequences efficiently - top_logprobs: 20 # Required for confidence computation - - # Filtering and voting - filter_method: "top10" # Use top-10 traces by confidence for voting - n_threads: 1 # Not used for vLLM (single batched call) - -# Dataset configuration -dataset: - subset: 30 # All AIME 2025 problems - offset: 0 # Start from beginning - prompt_file: "${hydra:runtime.cwd}/config/prompts/default.txt" diff --git a/config/experiments/deepconf/deepconf_vllm_qwen3_aime2025_nothink.yaml b/config/experiments/deepconf/deepconf_vllm_qwen3_aime2025_nothink.yaml deleted file mode 100644 index e42ce523..00000000 --- a/config/experiments/deepconf/deepconf_vllm_qwen3_aime2025_nothink.yaml +++ /dev/null @@ -1,74 +0,0 @@ -# @package _global_ -# Run: CUDA_VISIBLE_DEVICES=0 python scripts/run_tts_eval.py --config-path=../config --config-name=experiments/deepconf/deepconf_vllm_qwen3_aime2025_nothink -# Note: Do NOT specify CUDA_VISIBLE_DEVICES when submitting via scripts/slurm/submit.sh - SLURM assigns GPUs automatically -# DeepConf Offline Mode: Qwen3-8B via vLLM on AIME 2025 (No Thinking) -# Fast inference with PagedAttention - no OOM on long sequences - -defaults: - - /config - - /dataset/aime_2025 - - /model/vllm_qwen3 - - /generation/default - - /system/default - - /evaluation/default - - _self_ - -# Run naming -run_name: "deepconf_vllm_qwen3_aime2025_nothink_seed${system.seed}_${now:%H-%M-%S}" - -# Main configuration -verbose: true -report_to: null # Disable wandb for local runs - -# Model configuration (vLLM backend) -model: - type: "vllm" - model_path: "Qwen/Qwen3-8B" - model_short_name: "qwen3_8b" - gpu_memory_utilization: 0.9 - tensor_parallel_size: 1 - enable_prefix_caching: true - trust_remote_code: true - max_context_budget: 32768 - disable_thinking_mode: true # Disable Qwen3 thinking mode - -# System configuration -system: - device: cuda - seed: 42 - -# Generation configuration (non-thinking mode: temp=0.7, top_p=0.8, top_k=20) -# See: https://huggingface.co/Qwen/Qwen3-8B -generation: - max_new_tokens: 32768 # AIME problems require more tokens - temperature: 0.7 - top_p: 0.8 - top_k: 20 - batch_size: 1 - max_length: 8192 - -# No scorer needed (DeepConf has built-in weighted voting) -scorer: null - -# DeepConf offline mode strategy -strategy: - type: deepconf - mode: "offline" # Generate all traces, filter by confidence, weighted vote - budget: 16 # Number of reasoning traces to generate (all in ONE vLLM call!) - - # Generation parameters (non-thinking mode) - window_size: 2048 # Sliding window for confidence computation - temperature: 0.7 # Non-thinking mode standard - top_p: 0.8 # Non-thinking mode standard - max_tokens: 32768 # vLLM handles long sequences efficiently - top_logprobs: 20 # Required for confidence computation - - # Filtering and voting - filter_method: "top10" # Use top-10 traces by confidence for voting - n_threads: 1 # Not used for vLLM (single batched call) - -# Dataset configuration -dataset: - subset: 30 # All AIME 2025 problems - offset: 0 # Start from beginning - prompt_file: "${hydra:runtime.cwd}/config/prompts/default.txt" diff --git a/config/experiments/deepconf/deepconf_vllm_qwen3_math500.yaml b/config/experiments/deepconf/deepconf_vllm_qwen3_math500.yaml deleted file mode 100644 index ea2fa63c..00000000 --- a/config/experiments/deepconf/deepconf_vllm_qwen3_math500.yaml +++ /dev/null @@ -1,74 +0,0 @@ -# @package _global_ -# Run: CUDA_VISIBLE_DEVICES=0 python scripts/run_tts_eval.py --config-path=../config --config-name=experiments/deepconf/deepconf_vllm_qwen3_math500 -# Note: Do NOT specify CUDA_VISIBLE_DEVICES when submitting via scripts/slurm/submit.sh - SLURM assigns GPUs automatically -# DeepConf Offline Mode: Qwen3-8B via vLLM on MATH-500 - -defaults: - - /config - - /dataset/math_500 - - /model/vllm_qwen3 - - /generation/default - - /system/default - - /evaluation/default - - _self_ - -# Run naming -run_name: "deepconf_vllm_qwen3_math500_seed${system.seed}_${now:%H-%M-%S}" - -# Main configuration -verbose: true -report_to: null # Disable wandb for local runs - -# Model configuration (vLLM backend) -model: - type: "vllm" - model_path: "Qwen/Qwen3-8B" - model_short_name: "qwen3_8b" - gpu_memory_utilization: 0.9 - tensor_parallel_size: 1 - enable_prefix_caching: true - trust_remote_code: true - max_context_budget: 32768 - disable_thinking_mode: true # Disable Qwen3 thinking mode - -# System configuration -system: - device: cuda - seed: 42 - -# Generation configuration (non-thinking mode: temp=0.7, top_p=0.8, top_k=20) -generation: - max_new_tokens: 32768 - temperature: 0.7 - top_p: 0.8 - top_k: 20 - batch_size: 1 - max_length: 8192 - -# No scorer needed (DeepConf has built-in weighted voting) -scorer: null - -# DeepConf offline mode strategy -strategy: - type: deepconf - mode: "offline" # Generate all traces, filter by confidence, weighted vote - budget: 16 # Number of reasoning traces to generate - - # Generation parameters - window_size: 2048 # Sliding window for confidence computation - temperature: 0.7 # Sampling temperature for diversity - top_p: 0.8 # Nucleus sampling - max_tokens: 32768 # vLLM handles long sequences efficiently - top_logprobs: 20 # Required for confidence computation - - # Filtering and voting - filter_method: "top10" # Use top-10 traces by confidence for voting - n_threads: 1 # Not used for vLLM (single batched call) - -# Dataset configuration -dataset: - subset: 500 # All MATH-500 problems - offset: 0 - question_field: "problem" - answer_field: "answer" - prompt_file: "${hydra:runtime.cwd}/config/prompts/default_relaxed.txt" diff --git a/config/experiments/deepconf/run_aime_2025_deepconf_offline.yaml b/config/experiments/deepconf/run_aime_2025_deepconf_offline.yaml deleted file mode 100644 index e816d4b5..00000000 --- a/config/experiments/deepconf/run_aime_2025_deepconf_offline.yaml +++ /dev/null @@ -1,74 +0,0 @@ -# @package _global_ -# Run: CUDA_VISIBLE_DEVICES=0 python scripts/run_tts_eval.py --config-path=../config --config-name=experiments/deepconf/run_aime_2025_deepconf_offline -# Note: Do NOT specify CUDA_VISIBLE_DEVICES when submitting via scripts/slurm/submit.sh - SLURM assigns GPUs automatically - -# DeepConf Offline Mode Configuration for AIME 2025 -# Competition-level mathematics problems from American Invitational Mathematics Examination - -defaults: - - /config # Load main Hydra config - - /dataset/aime_2025 - - /model/openrouter - - /generation/default - - /system/default - - _self_ - -# Run naming for descriptive outputs -run_name: "run_aime_2025_deepconf_offline_seed${system.seed}_${now:%H-%M-%S}" - -# Main configuration -verbose: false -report_to: wandb # Enable wandb logging (optional) -wandb_project: llm-tts-eval-deepconf # Wandb project name for DeepConf experiments -wandb_group: "deepconf_aime2025" - -# Model configuration (API-based) -model: - provider: openrouter - model_name: "openai/gpt-4o-mini" - model_short_name: "openrouter_gpt4o_mini" - api_key: null # Set via OPENROUTER_API_KEY env var - top_logprobs: 20 # Required for DeepConf - -# Generation settings -generation: - max_new_tokens: 32000 # Large limit for complex AIME problems - temperature: 0.7 - top_p: 1.0 - batch_size: 1 - -# DeepConf strategy (no scorer needed) -scorer: null - -# DeepConf offline mode strategy -strategy: - type: deepconf - mode: "offline" # Standard weighted voting mode - budget: 16 # Number of traces to generate - - # Generation parameters - window_size: 2048 # Sliding window (algorithmic parameter - DO NOT CHANGE) - temperature: 0.6 # Match original wrapper.py default - top_p: 0.95 # Match original wrapper.py default - max_tokens: 32000 # Match original wrapper.py default - top_logprobs: 20 - - # Filtering and voting - filter_method: "top10" # Use top-10 traces by confidence - -# Evaluation configuration -evaluation: - evaluators: - - llm_judge - llm_judge: - provider: openrouter - base_url: https://openrouter.ai/api/v1 - model: deepseek/deepseek-r1-0528 - cache_path: ~/.cache - n_threads: 1 - prompt_file: null - -# Dataset settings -dataset: - subset: null # Evaluate all AIME 2025 problems (30 total) - prompt_file: "./config/prompts/aime_2025_minimal.txt" # Minimal zero-shot prompt (matching original DeepConf repo) diff --git a/config/experiments/deepconf/run_gsm8k_deepconf_offline.yaml b/config/experiments/deepconf/run_gsm8k_deepconf_offline.yaml deleted file mode 100644 index 990e1e4f..00000000 --- a/config/experiments/deepconf/run_gsm8k_deepconf_offline.yaml +++ /dev/null @@ -1,73 +0,0 @@ -# @package _global_ -# Run: CUDA_VISIBLE_DEVICES=0 python scripts/run_tts_eval.py --config-path=../config --config-name=experiments/deepconf/run_gsm8k_deepconf_offline -# Note: Do NOT specify CUDA_VISIBLE_DEVICES when submitting via scripts/slurm/submit.sh - SLURM assigns GPUs automatically - -# DeepConf Offline Mode Configuration -# This demonstrates standard weighted voting without early stopping - -defaults: - - /config # Load main Hydra config - - /dataset/gsm8k - - /model/openrouter - - /generation/default - - /system/default - - _self_ - -# Run naming for descriptive outputs -run_name: "run_gsm8k_deepconf_offline_seed${system.seed}_${now:%H-%M-%S}" - -# Main configuration -verbose: false -report_to: wandb # Enable wandb logging (optional) -wandb_project: llm-tts-eval-deepconf # Wandb project name for DeepConf experiments -wandb_group: "deepconf_gsm8k" - -# Model configuration (API-based) -model: - provider: openrouter - model_name: "openai/gpt-4o-mini" - model_short_name: "openrouter_gpt4o_mini" - api_key: null # Set via OPENROUTER_API_KEY env var - top_logprobs: 20 # Required for DeepConf - -# Generation settings -generation: - max_new_tokens: 512 - temperature: 0.7 - top_p: 1.0 - batch_size: 1 - -# DeepConf strategy (no scorer needed) -scorer: null - -# DeepConf offline mode strategy -strategy: - type: deepconf - mode: "offline" # Standard weighted voting mode - budget: 8 # Number of traces to generate - - # Generation parameters - window_size: 2048 # Sliding window (algorithmic parameter - DO NOT CHANGE) - temperature: 0.6 # Match original wrapper.py default - top_p: 0.95 # Match original wrapper.py default - max_tokens: 32000 # Match original wrapper.py default - top_logprobs: 20 - - # Filtering and voting - filter_method: "top5" # Use top-5 traces by confidence - -# Evaluation configuration -evaluation: - evaluators: - - llm_judge - llm_judge: - provider: openrouter - base_url: https://openrouter.ai/api/v1 - model: deepseek/deepseek-r1-0528 - cache_path: ~/.cache - n_threads: 1 - prompt_file: null - -# Dataset settings -dataset: - subset: 10 # Number of examples to evaluate diff --git a/config/experiments/deepconf/run_gsm8k_deepconf_online.yaml b/config/experiments/deepconf/run_gsm8k_deepconf_online.yaml deleted file mode 100644 index 0a297fa7..00000000 --- a/config/experiments/deepconf/run_gsm8k_deepconf_online.yaml +++ /dev/null @@ -1,77 +0,0 @@ -# @package _global_ -# Run: CUDA_VISIBLE_DEVICES=0 python scripts/run_tts_eval.py --config-path=../config --config-name=experiments/deepconf/run_gsm8k_deepconf_online -# Note: Do NOT specify CUDA_VISIBLE_DEVICES when submitting via scripts/slurm/submit.sh - SLURM assigns GPUs automatically - -# DeepConf Online Mode Configuration -# This demonstrates adaptive early stopping with streaming + logprobs - -# Run naming -run_name: "seed${system.seed}_${now:%H-%M-%S}" - -defaults: - - /dataset/gsm8k - - /model/openrouter - - /generation/default - - /output/default - - /system/default - - _self_ - -# Main configuration -verbose: false - -# Model configuration (API-based) -model: - provider: openrouter - model_name: "openai/gpt-4o-mini" - model_short_name: "openrouter_gpt4o_mini" - api_key: null # Set via OPENROUTER_API_KEY env var - top_logprobs: 20 # Required for DeepConf - -# Generation settings -generation: - max_new_tokens: 512 - temperature: 0.7 - top_p: 1.0 - batch_size: 1 - -# DeepConf strategy (no scorer needed) -scorer: null - -# DeepConf online mode strategy -strategy: - type: deepconf - mode: "online" # Enable online mode with adaptive early stopping - - # Budget (not used in online mode but required by config) - budget: 10 - - # Online mode parameters - warmup_traces: 4 # Initial traces to establish confidence threshold - total_budget: 10 # Total traces (warmup + adaptive) - confidence_percentile: 90 # Use 90th percentile (algorithmic parameter) - - # Generation parameters - window_size: 2048 # Sliding window (algorithmic parameter - DO NOT CHANGE) - temperature: 0.7 - top_p: 1.0 - max_tokens: 512 - top_logprobs: 20 - - # Filtering and voting - filter_method: "top5" # Use top-5 traces by confidence - -# Evaluation configuration -evaluation: - evaluators: - - llm_judge - llm_judge: - provider: openrouter - base_url: https://openrouter.ai/api/v1 - model: openai/gpt-3.5-turbo - cache_path: ~/.cache - n_threads: 1 - prompt_file: null - -# Dataset settings -dataset: - subset: 10 # Number of examples to evaluate diff --git a/config/strategy/deepconf.yaml b/config/strategy/deepconf.yaml index 59b52370..e8796f6a 100644 --- a/config/strategy/deepconf.yaml +++ b/config/strategy/deepconf.yaml @@ -6,3 +6,6 @@ budget: 8 # Number of reasoning traces to generate window_size: 2048 # Sliding window size for confidence computation filter_method: "top10" # Options: "none", "top10", "threshold" confidence_threshold: null # Only used if filter_method is "threshold" +mode: "offline" # "offline" (single batch) or "online" (incremental with early stopping) +online_batch_size: 4 # Traces per batch in online mode +min_agreement: 0.8 # Agreement threshold for early stopping in online mode diff --git a/docs/deepconf/DeepConf.md b/docs/deepconf/DeepConf.md index c923977e..f1f1d052 100644 --- a/docs/deepconf/DeepConf.md +++ b/docs/deepconf/DeepConf.md @@ -1,6 +1,6 @@ # DeepConf - Deep Think with Confidence -Complete guide for DeepConf implementation with OpenRouter API. +Guide for the DeepConf strategy integrated with the framework's `step_generator`. **Paper**: [Deep Think with Confidence](https://arxiv.org/abs/2508.15260) **Original Code**: [github.com/facebookresearch/deepconf](https://github.com/facebookresearch/deepconf) @@ -11,457 +11,344 @@ Complete guide for DeepConf implementation with OpenRouter API. 1. [Quick Start](#quick-start) 2. [Algorithm Overview](#algorithm-overview) -3. [Online vs Offline Mode](#online-vs-offline-mode) -4. [Implementation Details](#implementation-details) -5. [Configuration](#configuration) -6. [Testing](#testing) -7. [Troubleshooting](#troubleshooting) +3. [Implementation Details](#implementation-details) +4. [Configuration](#configuration) +5. [Troubleshooting](#troubleshooting) --- ## Quick Start -### Installation +### Running an Experiment ```bash -pip install openai -export OPENROUTER_API_KEY="your-key" +python scripts/run_tts_eval.py \ + --config-path ../config \ + --config-name experiments/deepconf/your_config \ + dataset.subset=10 ``` -### Basic Usage - -```python -from llm_tts.models import BlackboxModelWithStreaming -from llm_tts.strategies.deepconf import StrategyDeepConf - -# Create model -model = BlackboxModelWithStreaming( - openai_api_key="your-key", - model_path="openai/gpt-4o-mini", - supports_logprobs=True, - base_url="https://openrouter.ai/api/v1" -) - -# Create strategy -strategy = StrategyDeepConf( - model=model, - budget=8, # Number of reasoning traces - window_size=2048, # Sliding window size - filter_method="top10" # Keep top 10% by confidence -) - -# Run DeepConf -prompt = "What is 23 + 47? Put your answer in \\boxed{}." -result = strategy.generate_trajectory(prompt) - -print(f"Answer: {result['metadata']['selected_answer']}") -print(f"Confidence: {result['metadata']['confidence_score']:.2%}") -``` +### Minimal Config -### Running Experiments +Create a config YAML under `config/experiments/deepconf/`: -**Offline mode** (generate all traces, then filter & vote): -```bash -python scripts/run_tts_eval.py \ - --config-path ../config \ - --config-name experiments/deepconf/run_gsm8k_deepconf \ - dataset.subset=10 \ - strategy.mode=offline \ - strategy.budget=8 -``` +```yaml +# @package _global_ +defaults: + - /config + - /dataset/math_500 + - /model/vllm_qwen3 + - /generation/default + - /system/default + - /evaluation/default + - _self_ -**Online mode** (warmup phase + adaptive generation): -```bash -python scripts/run_tts_eval.py \ - --config-path ../config \ - --config-name experiments/deepconf/run_gsm8k_deepconf \ - dataset.subset=10 \ - strategy.mode=online \ - strategy.warmup_traces=4 \ - strategy.total_budget=16 +run_name: "deepconf_example_${now:%H-%M-%S}" + +scorer: null # DeepConf has built-in weighted voting + +strategy: + type: deepconf + min_step_tokens: 50 + budget: 16 + window_size: 2048 + filter_method: "top10" ``` -**Key parameters:** -- `strategy.mode`: `offline` or `online` -- `strategy.budget`: Number of traces (offline mode) -- `strategy.warmup_traces`: Warmup traces (online mode) -- `strategy.total_budget`: Total traces including warmup (online mode) -- `strategy.filter_method`: `top10`, `top5`, or `threshold` -- `strategy.window_size`: Sliding window for confidence (default: 2048) -- `eval_method`: `simple` (numeric comparison) or `llm_judge` (LLM-based verification) +### Key Parameters + +| Parameter | Description | Default | +|-----------|-------------|---------| +| `strategy.budget` | Number of traces to generate per sample | 8 | +| `strategy.window_size` | Sliding window size for confidence | 2048 | +| `strategy.filter_method` | `"top10"`, `"top5"`, `"none"`, or `"threshold"` | `"top10"` | +| `strategy.confidence_threshold` | Threshold value (only when `filter_method="threshold"`) | null | +| `strategy.data_name` | Dataset name for answer extraction | null | +| `strategy.mode` | `"offline"` (single batch) or `"online"` (incremental) | `"offline"` | +| `strategy.online_batch_size` | Traces per batch in online mode | 4 | +| `strategy.min_agreement` | Agreement threshold for early stopping (online) | 0.8 | +| `generation.temperature` | Sampling temperature (controls diversity) | 0.7 | +| `generation.max_new_tokens` | Max tokens per trace | 32768 | + +Generation parameters (`temperature`, `top_p`, `max_new_tokens`) are configured in the `generation` section and picked up by the step_generator automatically. --- ## Algorithm Overview -### Core Idea +### Modes -DeepConf generates multiple reasoning traces, scores them by confidence using token probabilities, filters low-confidence traces, and performs weighted voting. +DeepConf supports two execution modes: + +- **Offline** (default): Generate all N traces in a single batched call, then filter and vote. Most efficient when you want a fixed compute budget. +- **Online**: Generate traces in small batches, check consensus after each batch, and stop early when filtered traces agree. Saves compute when traces converge early. + +### Core Idea (Offline) + +DeepConf generates multiple reasoning traces in a single batched call, scores each trace by confidence using token-level logprobs, filters low-confidence traces, and performs weighted majority voting. ### Algorithm Steps #### 1. Generate Multiple Traces -```python -for i in range(budget): - trace_i = model.generate(prompt, temperature=0.7) -``` + +All N traces are generated in a **single** `generate_step_candidates_batch()` call via the framework's step_generator (vLLM or API backend). This is efficient because vLLM can batch all traces with shared prefix caching. #### 2. Calculate Token Confidence -**Formula**: `C_i = -(1/k) * Σ_{j=1}^k log P_i(j)` + +For each token position, extract top-k logprobs from `StepCandidate.other_data["raw_logprobs"]`: + +**Formula**: `C_i = -(1/k) * sum_{j=1}^{k} log P_i(j)` Where: - `C_i` = confidence for token i -- `k` = top-k token probabilities (typically 20) +- `k` = number of top logprob values (up to 20) - `P_i(j)` = probability of j-th most likely token -Higher score = more confident. +Higher score = more uncertain (lower confidence). #### 3. Sliding Window Group Confidence -**Formula**: `CG_i = (1/|G_i|) * Σ_{t∈G_i} C_t` + +**Formula**: `CG_i = (1/|G_i|) * sum_{t in G_i} C_t` Where: - `CG_i` = group confidence for window at position i - `G_i` = sliding window of size w starting at i -**IMPORTANT: Windows overlap!** -- Windows slide by 1 token -- Each window shares (window_size - 1) tokens with the next -- Source: `deepconf/deepconf/utils.py:59-60` +Windows overlap by sliding 1 token at a time. Each window shares `(window_size - 1)` tokens with the next. **Trace confidence** = `min(all_window_confidences)` #### 4. Filter Traces -- **top10**: Keep top 10% by confidence -- **threshold**: Keep traces above fixed threshold + +- **top10** / **top5**: Keep top K traces by confidence (sorted by min_conf descending) +- **threshold**: Keep traces above a fixed confidence threshold - **none**: Use all traces +If no traces pass the filter, all traces are used as fallback. + #### 5. Weighted Voting -```python + +``` for trace in filtered_traces: weight[trace.answer] += trace.min_conf selected = argmax(weight) ``` -Each trace "votes" with weight = its confidence. - -### Modes - -**Offline Mode** (Simpler): -1. Generate all N traces -2. Calculate confidence for each -3. Filter by confidence -4. Weighted voting - -**Online Mode** (Efficient - **NOW FULLY IMPLEMENTED**): -1. Warmup: Generate K traces, calibrate threshold -2. Adaptive: Stream tokens with logprobs, stop early when confidence drops -3. Filter and vote on high-confidence traces +Each trace "votes" with weight = its minimum window confidence. ---- +### Online Mode Algorithm -## Online vs Offline Mode - -### Offline Mode (Batch Generation) - -**How it works:** -1. Generate N complete traces with logprobs -2. Compute confidence for each complete trace -3. Filter traces by confidence threshold or top-K -4. Weighted majority voting - -**Advantages:** -- Simple implementation -- More stable confidence measurements (full traces) -- Works with any API that supports logprobs - -**Disadvantages:** -- Always generates full traces (wasteful if low confidence) -- Higher token cost -- Slower execution - -**Use when:** -- You want maximum accuracy -- Token cost is not a concern -- You need stable, reproducible results - -### Online Mode (Adaptive Early Stopping) - -**How it works:** -1. **Warmup Phase**: Generate K traces to calibrate confidence threshold - - Example: K=3 traces, use 90th percentile of min confidences -2. **Adaptive Phase**: Generate remaining (N-K) traces with early stopping - - Stream tokens one-by-one with logprobs - - Maintain sliding window of token confidences - - Stop generation when window confidence < threshold -3. **Filter & Vote**: Same as offline mode - -**Technical Implementation:** -- Uses `stream=True` + `logprobs=True` (OpenRouter/Together AI support this!) -- `ConfidenceProcessor` maintains sliding window and triggers early stop -- `BlackboxModelWithStreaming.generate_texts()` with `confidence_callback` handles streaming -- Both offline and online modes use the same streaming path with logprobs - -**Advantages:** -- **Token efficiency**: Stops generating low-confidence traces early -- **Faster**: Don't wait for full traces when confidence is low -- **Cost effective**: Can save 30-50% tokens vs offline mode - -**Disadvantages:** -- Slightly more complex implementation -- Warmup phase adds overhead for small budgets -- Early stopping might miss later recoveries - -**Use when:** -- Token cost matters -- You have budget >5 traces (warmup overhead worth it) -- You want faster inference - -### Configuration Examples - -**Offline Mode:** -```yaml -strategy: - mode: "offline" - budget: 10 # Generate 10 complete traces - filter_method: "top5" # Use top 5 by confidence - window_size: 5 - temperature: 0.7 - max_tokens: 512 -``` +In online mode, traces are generated incrementally and samples stop independently: -**Online Mode:** -```yaml -strategy: - mode: "online" - warmup_traces: 3 # Calibration phase - total_budget: 10 # 3 warmup + 7 adaptive - confidence_percentile: 10 # Use 90th percentile (100-10) - filter_method: "top5" - window_size: 5 - temperature: 0.7 - max_tokens: 512 ``` +all_traces = [[] for M samples] +completed = [False] * M -### Performance Comparison +while any sample not completed: + active = [samples not yet completed] -**Example: GSM8K with budget=10** + # Generate online_batch_size traces for each active sample + candidates = step_generator.generate_step_candidates_batch( + requests=[requests[i] for i in active], + candidates_per_step=online_batch_size, + sample_ids=active, + ) -| Mode | Avg Tokens/Question | Time | Accuracy | -|------|---------------------|------|----------| -| Offline | ~5000 | 45s | 85% | -| Online (p=10) | ~3500 | 35s | 84% | -| Online (p=30) | ~2800 | 28s | 82% | + for each active sample: + # Compute confidence, extract answers for new traces + # Add new traces to pool + # Check: budget exhausted → mark completed + # Check: agreement >= min_agreement → early stop +``` -*Note: Online mode with higher confidence percentile (stricter threshold) stops earlier but may sacrifice some accuracy.* +Key properties: +- **Per-sample tracking**: Samples that converge early stop generating, while others continue. This saves compute on "easy" problems. +- **Prefix caching**: vLLM's prefix caching works across batch calls for the same prompt, so incremental generation is efficient. +- **Budget is a hard cap**: Online mode never exceeds `budget` traces per sample, even if consensus is not reached. +- **Same output format**: Online and offline produce identical result dictionaries. Online adds an `early_stopped` flag. --- ## Implementation Details -### Key Files +### Architecture + +DeepConf uses the framework's `step_generator` pattern (same as `StrategySelfConsistency`): ``` -llm_tts/ -├── models/ -│ ├── blackboxmodel_with_streaming.py # Unified streaming with logprobs support -│ └── base.py # Base model interface -├── strategies/ -│ └── deepconf/ # DeepConf implementation -│ ├── strategy.py # Main strategy (offline & online modes) -│ └── utils.py # Confidence computation utilities -├── early_stopping.py # Early stopping conditions -├── step_boundary_detector.py # Detects step/answer boundaries -└── scorers/ - └── majority_voting.py # Weighted voting implementation - -config/experiments/deepconf/ -├── run_gsm8k_deepconf_offline.yaml # Offline mode config -└── run_gsm8k_deepconf_online.yaml # Online mode config - -tests/deepconf/ -├── test_deepconf_accurate.py # Unit tests -├── test_online_mode.py # Online mode tests -└── test_deepconf_math.py # Math validation tests +StrategyDeepConf + └── step_generator.generate_step_candidates_batch() + ├── VLLMStepGenerator (vLLM backend) + └── StepCandidateGeneratorThroughAPI (API backend) ``` -### Streaming + Logprobs Architecture - -**Key Insight**: Both OpenRouter and Together AI support `stream=True` + `logprobs=True` simultaneously! +### Key Files -**Implementation** (`BlackboxModelWithStreaming.generate_texts()`): -```python -# Single unified path for both offline and online modes -if args.get("output_scores", False) or args.get("stream_with_confidence", False): - response = self.client.chat.completions.create( - model=self.model_path, - messages=chat, - stream=True, # Stream tokens - logprobs=True, # Get logprobs for each token - top_logprobs=20, # Top-20 for confidence calculation - ) +``` +llm_tts/strategies/deepconf/ +├── __init__.py # Package exports +├── strategy.py # Main strategy (~475 lines) +└── utils.py # compute_sliding_window_confidence() - # Process streaming chunks - for chunk in response: - # Extract token + logprobs - # Call confidence_callback if provided (online mode) - # Accumulate text + logprobs +config/strategy/ +└── deepconf.yaml # Base strategy config ``` -**Offline Mode**: Calls `generate_texts()` with `output_scores=True`, no callback -**Online Mode**: Calls `generate_texts()` with `confidence_callback` for early stopping +### Generation Flow -### Answer Extraction +1. `generate_trajectories_batch(requests, sample_indices)` is called with M samples +2. Single `step_generator.generate_step_candidates_batch()` call generates M x N candidates +3. For thinking mode: `_complete_thinking_paths()` generates answer phases for candidates that stopped at `` +4. For each candidate: extract `raw_logprobs` from `StepCandidate.other_data` and compute confidence via `_compute_trace_confidence()` +5. Per-sample `_filter_and_vote()` selects the best answer -DeepConf expects **LaTeX boxed format**: -``` -\boxed{answer} -``` +### Logprob Format + +The `raw_logprobs` from vLLM is `List[Dict[token_id -> Logprob]]` where each `Logprob` object has a `.logprob` attribute. For API models, `APILogprobData` objects provide the same interface. The confidence computation works identically for both backends. + +### Thinking Mode Support -**Examples:** -- `\boxed{70}` → `"70"` ✅ -- `\boxed{x+y}` → `"x+y"` ✅ -- Plain text → `None` ❌ +When `thinking_mode=True` (set via `model.disable_thinking_mode: false`): +- Generation stops at `` (added to stop tokens) +- Answer phase is generated via `generate_answer_candidates_batch()` +- Full text = thinking + answer concatenated +- Confidence is computed from the thinking-phase logprobs -### Key Improvements Over Previous Version +### Answer Extraction -| Issue | Before | After | -|-------|--------|-------| -| Answer extraction | "the final answer is 70" ❌ | "70" ✅ | -| Confidence | Heuristic text-based | Real token logprobs ✅ | -| Formula | Approximations | Exact DeepConf formula ✅ | -| Windows | Not clear | Confirmed overlapping ✅ | +Uses the framework's `llm_tts.utils.extract_answer()` with `answer_format="auto"`: +- Tries `: ... ` format first +- Falls back to `\boxed{...}` format +- Supports nested braces --- ## Configuration -### Strategy Parameters - -```python -StrategyDeepConf( - model=model, - budget=8, # Number of traces (8-512) - window_size=2048, # Sliding window (16, 32, 2048) - temperature=0.7, # Sampling temperature (0.6-0.8) - top_p=0.95, # Top-p sampling - max_tokens=4096, # Max tokens per trace - top_logprobs=20, # Top-k for confidence - filter_method="top10" # none, top10, threshold -) +### Base Config (`config/strategy/deepconf.yaml`) + +```yaml +type: deepconf +min_step_tokens: 50 +budget: 8 +window_size: 2048 +filter_method: "top10" +confidence_threshold: null +mode: "offline" +online_batch_size: 4 +min_agreement: 0.8 ``` ### Filter Methods -1. **`"none"`** - No filtering, use all traces -2. **`"top10"`** - Keep top 10% by confidence -3. **`"threshold"`** - Keep traces above `confidence_threshold` +| Method | Behavior | +|--------|----------| +| `"none"` | No filtering, use all traces | +| `"top10"` | Keep top 10 traces by confidence | +| `"top5"` | Keep top 5 traces by confidence | +| `"threshold"` | Keep traces with `min_conf >= confidence_threshold` | -### Example Configurations +### Example: vLLM with Qwen3 -**High Accuracy**: -```python -strategy = StrategyDeepConf(model, budget=16, filter_method="top10") -``` +```yaml +model: + type: "vllm" + model_path: "Qwen/Qwen3-8B" + disable_thinking_mode: true # Non-thinking mode + gpu_memory_utilization: 0.9 + enable_prefix_caching: true + +generation: + max_new_tokens: 32768 + temperature: 0.7 + top_p: 0.8 + top_k: 20 -**Fast Prototyping**: -```python -strategy = StrategyDeepConf(model, budget=5, filter_method="none") -``` +scorer: null -**Custom Threshold**: -```python -strategy = StrategyDeepConf( - model, - budget=8, - filter_method="threshold", - confidence_threshold=12.0 -) +strategy: + type: deepconf + min_step_tokens: 50 + budget: 16 + window_size: 2048 + filter_method: "top10" ``` ---- +### Example: High Budget -## Testing +```yaml +strategy: + type: deepconf + budget: 64 + filter_method: "top10" + window_size: 2048 +``` -### Integration Tests +### Example: Custom Threshold -```bash -python tests/deepconf/test_deepconf_accurate.py +```yaml +strategy: + type: deepconf + budget: 16 + filter_method: "threshold" + confidence_threshold: 12.0 + window_size: 2048 ``` -Tests: -1. ✅ Answer Extraction - LaTeX parsing -2. ✅ Model Logprobs Support -3. ✅ DeepConf Generation -4. ✅ Weighted Voting - -### Math Problems +### Example: Online Mode -```bash -# Standard run -python tests/deepconf/test_deepconf_math.py +```yaml +strategy: + type: deepconf + budget: 16 + mode: "online" + online_batch_size: 4 + min_agreement: 0.8 + window_size: 2048 + filter_method: "top10" +``` -# Verbose - see full reasoning paths -python tests/deepconf/test_deepconf_math.py --verbose +With online mode, easy samples may stop after just 4 traces (one batch) if all agree, while hard samples continue up to the full budget of 16. -# More traces -python tests/deepconf/test_deepconf_math.py --budget 10 -``` +### When to Use Online vs Offline -Problems: -- Difference of squares: 15² - 8² -- Rectangle area: 12 × 8 -- Speed-distance: 120 km in 2 hours -- Arithmetic series: 1+2+...+10 -- Order of operations: (3+4) × (5+6) +| Scenario | Recommended Mode | +|----------|-----------------| +| Fixed compute budget, all samples similar difficulty | Offline | +| Mixed difficulty, want to save compute on easy samples | Online | +| Small budget (≤ 8) | Offline (overhead not worth it) | +| Large budget (≥ 16) with diverse difficulty | Online | +| Benchmarking/reproducing paper results | Offline | --- ## Troubleshooting -### "Model does not support logprobs" - -**Solution**: Use OpenRouter with OpenAI models: -- ✅ gpt-4o-mini -- ✅ gpt-4o -- ✅ gpt-3.5-turbo - ### "No valid answers extracted" -**Solution**: Include `\boxed{}` in prompt: -```python -prompt = "Calculate X. Put answer in \\boxed{}." +Ensure the prompt instructs the model to use `\boxed{}` format: ``` - -### "All traces have same answer" - -**Solution**: Increase temperature for diversity: -```python -strategy = StrategyDeepConf(model, temperature=0.9) +Put your final answer in \boxed{}. ``` -### Import Errors +### "All traces have same answer" -**Solution**: Run from project root: -```bash -cd /path/to/llm-tts-service -python tests/deepconf/test_deepconf_accurate.py +Increase temperature for more diversity: +```yaml +generation: + temperature: 0.9 ``` ---- - -## Performance - -### OpenRouter + gpt-4o-mini +### All min_conf values are 0.0 or inf -- Budget=8, max_tokens=500: ~15-20s, ~$0.02 -- Budget=16, max_tokens=1000: ~40-50s, ~$0.08 +This means `raw_logprobs` is empty. Check: +- vLLM models always provide logprobs (hardcoded `logprobs=20`) +- API models need `supports_logprobs: true` in model config -### Cost Optimization +### Low accuracy despite high budget -1. Reduce `budget` (fewer traces) -2. Reduce `max_tokens` (shorter reasoning) -3. Use `filter_method="top10"` (quality over quantity) +Try adjusting the filter method. `"top10"` keeps only the 10 most confident traces for voting. If budget < 10, all traces pass the filter anyway. Consider `"top5"` for smaller budgets or `"none"` to use all traces. --- @@ -469,4 +356,3 @@ python tests/deepconf/test_deepconf_accurate.py - **Paper**: [arxiv.org/abs/2508.15260](https://arxiv.org/abs/2508.15260) - **Original Code**: [github.com/facebookresearch/deepconf](https://github.com/facebookresearch/deepconf) -- **OpenRouter**: [openrouter.ai/docs](https://openrouter.ai/docs) diff --git a/llm_tts/strategies/deepconf/__init__.py b/llm_tts/strategies/deepconf/__init__.py index 888ffbfc..97aa3fac 100644 --- a/llm_tts/strategies/deepconf/__init__.py +++ b/llm_tts/strategies/deepconf/__init__.py @@ -1,7 +1,7 @@ """ DeepConf strategy package. -Confidence-based test-time scaling with offline and online modes. +Confidence-based test-time scaling using framework's step_generator. """ from .strategy import StrategyDeepConf diff --git a/llm_tts/strategies/deepconf/strategy.py b/llm_tts/strategies/deepconf/strategy.py index 71943e0a..88b2b4b0 100644 --- a/llm_tts/strategies/deepconf/strategy.py +++ b/llm_tts/strategies/deepconf/strategy.py @@ -4,1222 +4,212 @@ Based on Facebook Research's DeepConf: https://github.com/facebookresearch/deepconf -Supports both offline and online modes: -- Offline: Generate N traces, filter by confidence, majority vote -- Online: Adaptive generation with confidence-based early stopping +Uses framework's step_generator for generation and extracts logprobs from +StepCandidate.other_data["raw_logprobs"] for confidence computation. + +Offline mode: + 1. Generate N complete traces via step_generator (single batched call) + 2. Compute per-token confidence from raw_logprobs + 3. Apply sliding window to get min_conf per trace + 4. Filter traces by confidence, majority vote on answers + +Online mode: + 1. Generate traces in small batches (online_batch_size per round) + 2. After each batch, check agreement among filtered traces + 3. Stop early when agreement >= min_agreement or budget exhausted + 4. Different samples converge independently """ import logging -import re -from typing import Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional -import numpy as np -import torch -from lm_polygraph import BlackboxModel -from transformers import StoppingCriteria, StoppingCriteriaList - -from llm_tts.early_stopping import ConfidenceEarlyStopping +from llm_tts.generators.base import convert_trajectory_to_string +from llm_tts.utils import extract_answer from ..metadata_builder import StrategyMetadataBuilder from ..strategy_base import StrategyBase -from .utils import ( - compute_sliding_window_confidence, - compute_token_confidence_from_logprobs, - extract_answer, -) - -log = logging.getLogger(__name__) +from .utils import compute_sliding_window_confidence +if TYPE_CHECKING: + from llm_tts.generators import StepCandidateGeneratorBase -class AnswerStoppingCriteria(StoppingCriteria): - """ - Stopping criteria that stops generation when a complete answer is found. +log = logging.getLogger(__name__) - Supports two formats: - - Default format: : ... - - Boxed format: \\boxed{...} - Works with batched generation (num_return_sequences > 1). +class StrategyDeepConf(StrategyBase): """ + DeepConf strategy using framework's step_generator. - def __init__( - self, - tokenizer, - start_length: int, - batch_size: int, - answer_format: str = "default", - min_tokens_after: int = 0, - ): - """ - Args: - tokenizer: HuggingFace tokenizer for decoding - start_length: Length of input tokens (to extract only generated text) - batch_size: Number of sequences being generated - answer_format: "default" for : format, "boxed" for \\boxed{} format - min_tokens_after: Minimum tokens to generate after finding answer marker - """ - self.tokenizer = tokenizer - self.start_length = start_length - self.batch_size = batch_size - self.answer_format = answer_format - self.min_tokens_after = min_tokens_after - self.finished = [False] * batch_size - # Track when we found answer for each sequence - self.answer_found_at = [None] * batch_size - - # Patterns for different formats - if answer_format == "default": - # Default format: : followed by or - self.end_pattern = re.compile(r"") - else: - # Boxed format - self.answer_pattern = re.compile(r"\\boxed\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}") - - def __call__( - self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs - ) -> bool: - """Check if all sequences have found their answer.""" + Generates N traces via generate_step_candidates_batch(), computes + per-token confidence from raw vLLM logprobs, filters by confidence, + and performs weighted majority voting. - for i in range(min(input_ids.shape[0], self.batch_size)): - if self.finished[i]: - continue - - # Decode only generated tokens - generated_ids = input_ids[i][self.start_length :] - generated_text = self.tokenizer.decode( - generated_ids, skip_special_tokens=True - ) - - # Check for answer pattern - if self.answer_format == "default": - # For default format, check for marker - if self.end_pattern.search(generated_text): - if self.answer_found_at[i] is None: - self.answer_found_at[i] = len(generated_ids) - tokens_after = len(generated_ids) - self.answer_found_at[i] - if tokens_after >= self.min_tokens_after: - self.finished[i] = True - else: - # For boxed format - if self.answer_pattern.search(generated_text): - if self.answer_found_at[i] is None: - self.answer_found_at[i] = len(generated_ids) - tokens_after = len(generated_ids) - self.answer_found_at[i] - if tokens_after >= self.min_tokens_after: - self.finished[i] = True - - # Stop when all sequences are finished - return all(self.finished) - - -class StrategyDeepConf(StrategyBase): - """ - DeepConf strategy with offline and online modes. - - Offline mode: - 1. Generate N complete traces with logprobs - 2. Compute confidence scores for each trace - 3. Filter traces by confidence threshold - 4. Majority vote on answers - - Online mode: - 1. Warmup: Generate K traces to estimate confidence threshold - 2. Adaptive: Generate remaining traces with early stopping - 3. Filter and majority vote + Supports two modes: + - offline: Generate all traces in a single batched call + - online: Generate traces in small batches with early stopping """ def __init__( self, - model, - mode: str, - budget: int, - window_size: int, - temperature: float, - top_p: float, - max_tokens: int, - top_logprobs: int, - filter_method: str, - warmup_traces: Optional[int] = None, - total_budget: Optional[int] = None, - confidence_percentile: Optional[int] = None, - confidence_threshold: Optional[float] = None, - n_threads: int = 8, - disable_thinking_mode: bool = True, - seed: int = 42, + step_generator: "StepCandidateGeneratorBase", + budget: int = 8, + window_size: int = 2048, + filter_method: str = "top10", + confidence_threshold: float = None, + data_name: str = None, + mode: str = "offline", + online_batch_size: int = 4, + min_agreement: float = 0.8, ): """ Initialize DeepConf strategy. Args: - model: Model supporting logprobs (BlackboxModel for API, or local HuggingFace model) - mode: "offline" or "online" - budget: Number of traces for offline mode - warmup_traces: Warmup traces for online mode - total_budget: Total budget for online mode - confidence_percentile: Confidence threshold percentile for online - window_size: Window size for confidence computation - temperature: Sampling temperature - top_p: Nucleus sampling parameter - max_tokens: Maximum tokens per trace - top_logprobs: Number of top logprobs to request - filter_method: Filtering method ("topK", "threshold", or specific like "top10") - confidence_threshold: Manual confidence threshold (optional) - n_threads: Number of threads for parallel API requests (default: 8) - disable_thinking_mode: Whether to disable thinking mode for models like Qwen3 - seed: Random seed for reproducibility (default: 42) + step_generator: Step generator for batched generation. + budget: Number of traces to generate per sample. + window_size: Sliding window size for confidence computation. + filter_method: Filtering method ("topK", "threshold", or "top10"). + confidence_threshold: Manual confidence threshold (for threshold filtering). + data_name: Dataset name for answer extraction. + mode: "offline" (single batch) or "online" (incremental with early stopping). + online_batch_size: Traces per batch in online mode. + min_agreement: Agreement threshold for early stopping in online mode. """ - self.model = model - self.mode = mode.lower() + self.step_generator = step_generator self.budget = budget - self.warmup_traces = warmup_traces - self.total_budget = total_budget - self.confidence_percentile = confidence_percentile self.window_size = window_size - self.temperature = temperature - self.top_p = top_p - self.max_tokens = max_tokens - self.top_logprobs = top_logprobs self.filter_method = filter_method self.confidence_threshold = confidence_threshold - self.n_threads = n_threads - self.disable_thinking_mode = disable_thinking_mode - self.seed = seed - - # Validate model supports logprobs (for API models) or has required attributes (for local) - if isinstance(model, BlackboxModel): - # API model: check supports_logprobs flag - if not hasattr(model, "supports_logprobs") or not model.supports_logprobs: - raise ValueError( - "API model must have supports_logprobs=True for DeepConf" - ) - else: - # Local HuggingFace model: check required attributes - required_attrs = ["model", "tokenizer", "tokenize", "device"] - missing = [attr for attr in required_attrs if not hasattr(model, attr)] - if missing: - raise ValueError( - f"Local model missing required attributes for DeepConf: {missing}" - ) + self.data_name = data_name + self.mode = mode + self.online_batch_size = online_batch_size + self.min_agreement = min_agreement - model_type = "API" if isinstance(model, BlackboxModel) else "Local" log.info( - f"DeepConf initialized ({model_type} model): mode={mode}, " - f"budget={budget if mode == 'offline' else total_budget}, " - f"window={window_size}, filter={filter_method}, threads={n_threads}" + f"DeepConf initialized: budget={budget}, window={window_size}, " + f"filter={filter_method}, mode={mode}" ) - - def generate_trajectory(self, prompt) -> Dict[str, Any]: - """ - Main entry point - generate traces and select best answer. - - Args: - prompt: Input prompt/question (str or list of messages) - - Returns: - Dictionary with trajectory, steps, and metadata - """ - # Handle both string and message list formats - if isinstance(prompt, list): - # Extract user message from message list - user_msg = next((m["content"] for m in prompt if m["role"] == "user"), None) - if user_msg is None: - raise ValueError("No user message found in prompt list") - prompt = user_msg - - if self.mode == "offline": - return self._generate_offline(prompt) - elif self.mode == "online": - return self._generate_online(prompt) - else: - raise ValueError(f"Invalid mode: {self.mode}") - - def _generate_offline(self, prompt: str) -> Dict[str, Any]: - """ - Offline mode: Generate all traces, then filter and vote. - """ - log.info(f"🎯 DeepConf Offline: Generating {self.budget} traces...") - - traces = self._generate_traces_batch(prompt, self.budget) - - log.info(f"Generated {len(traces)} traces") - for i, trace in enumerate(traces): + if mode == "online": log.info( - f" Trace {i+1}: min_conf={trace['min_conf']:.3f}, " - f"answer={trace['extracted_answer']}" + f" Online params: batch_size={online_batch_size}, " + f"min_agreement={min_agreement}" ) - # Filter and vote - result = self._filter_and_vote(traces) - - # Build enhanced metadata - filtered_set = set(id(t) for t in result["filtered_traces"]) - - # Full details for ALL traces (for analysis) - all_traces_details = [] - for i, trace in enumerate(traces): - all_traces_details.append( - { - "trace_id": i, - "text": trace["text"], # Full reasoning text - "min_conf": trace["min_conf"], - "mean_conf": ( - sum(trace["token_confs"]) / len(trace["token_confs"]) - if trace["token_confs"] - else 0.0 - ), - "answer": trace["extracted_answer"], - "num_tokens": len(trace.get("token_data", [])), - "selected": id(trace) in filtered_set, - "window_confs": trace.get( - "window_confs", [] - ), # For confidence charts - "token_confs": trace.get( - "token_confs", [] - ), # Per-token confidences - } - ) - - # Summary of filtered traces (for quick reference) - filtered_trace_ids = [ - traces.index(trace) for trace in result["filtered_traces"] - ] - - # Build metadata using StrategyMetadataBuilder - builder = StrategyMetadataBuilder("deepconf") - - # Add configuration - builder.add_config( - mode="offline", - budget=self.budget, - window_size=self.window_size, - temperature=self.temperature, - top_p=self.top_p, - max_tokens=self.max_tokens, - top_logprobs=self.top_logprobs, - filter_method=self.filter_method, - n_threads=self.n_threads, - ) - - # Add results - builder.add_results( - selected_answer=result["selected_answer"], - confidence_score=result["confidence_score"], - vote_distribution=result["vote_distribution"], - total_traces=len(traces), - filtered_traces=len(result["filtered_traces"]), - ) - - # Add generation details - builder.add_generation_details( - all_traces=all_traces_details, - filtered_trace_ids=filtered_trace_ids, - ) - - # Log summary to console - builder.log_summary(log) - - # Return all traces for analysis, with selected answer for evaluation - return { - "trajectory": result["selected_text"], # Winning trace text - "extracted_answer": result["selected_answer"], # For ExactMatch evaluation - "all_traces": all_traces_details, # All generated traces with full text - "steps": [t["text"] for t in all_traces_details], # All trace texts - "validity_scores": [ - t["min_conf"] for t in all_traces_details - ], # All confidences - "consensus_score": result[ - "agreement_rate" - ], # Proportion of traces with winning answer (0-1) - "vote_distribution": result["vote_distribution"], # Answer -> percentage - "completed": True, - "metadata": builder.build(), - } - - def _generate_online(self, prompt: str) -> Dict[str, Any]: - """ - Online mode: Warmup to get threshold, then adaptive generation. + def _complete_thinking_paths( + self, + request: List[Dict[str, str]], + candidates: List[Any], + ) -> List[Dict[str, Any]]: """ - log.info( - f"🎯 DeepConf Online: Warmup {self.warmup_traces}, " - f"Total {self.total_budget} traces..." - ) - - # Phase 1: Warmup - log.info("Phase 1: Warmup...") - - # Clear any existing early stopping for warmup phase - self.model.early_stopping = None + Complete thinking-mode candidates by generating answer phases. - warmup_traces = self._generate_traces_batch(prompt, self.warmup_traces) - - # Compute confidence threshold - warmup_min_confs = [t["min_conf"] for t in warmup_traces] - conf_threshold = float( - np.percentile(warmup_min_confs, 100 - self.confidence_percentile) - ) - - # Log warmup statistics - log.info(f"Warmup complete: conf_threshold={conf_threshold:.3f}") - log.info(f" Min confs: {[f'{c:.3f}' for c in warmup_min_confs]}") - log.info( - f" Mean: {np.mean(warmup_min_confs):.3f}, " - f"Median: {np.median(warmup_min_confs):.3f}, " - f"Std: {np.std(warmup_min_confs):.3f}" - ) - log.info( - f" Percentile: {100 - self.confidence_percentile}th = {conf_threshold:.3f}" - ) - - # Phase 2: Adaptive generation with early stopping - log.info("Phase 2: Adaptive generation...") - remaining = self.total_budget - self.warmup_traces - adaptive_traces = self._generate_traces_adaptive( - prompt, remaining, conf_threshold - ) - - # Calculate adaptive phase statistics - num_stopped_early = sum( - 1 for t in adaptive_traces if t.get("stopped_early", False) - ) - num_completed = len(adaptive_traces) - num_stopped_early - adaptive_tokens = sum(len(t.get("token_data", [])) for t in adaptive_traces) - warmup_tokens = sum(len(t.get("token_data", [])) for t in warmup_traces) - - log.info(f"Adaptive complete: generated {len(adaptive_traces)} traces") - log.info(f" Stopped early: {num_stopped_early}/{len(adaptive_traces)}") - log.info(f" Completed: {num_completed}/{len(adaptive_traces)}") - log.info(f" Tokens: warmup={warmup_tokens}, adaptive={adaptive_tokens}") - - if warmup_traces and adaptive_traces: - avg_warmup_tokens = warmup_tokens / len(warmup_traces) - avg_adaptive_tokens = adaptive_tokens / len(adaptive_traces) - if num_stopped_early > 0: - early_stopped = [ - t for t in adaptive_traces if t.get("stopped_early", False) - ] - avg_early_stop_tokens = ( - sum(len(t.get("token_data", [])) for t in early_stopped) - / num_stopped_early - ) - potential_tokens = num_stopped_early * avg_warmup_tokens - saved_tokens = potential_tokens - sum( - len(t.get("token_data", [])) for t in early_stopped - ) - savings_pct = ( - (saved_tokens / potential_tokens * 100) - if potential_tokens > 0 - else 0 - ) - log.info(f" Avg tokens: warmup={avg_warmup_tokens:.1f}, \ - adaptive={avg_adaptive_tokens:.1f}") - log.info(f" Early stop avg: {avg_early_stop_tokens:.1f} tokens") - log.info(f" Token savings: {saved_tokens:.0f} ({savings_pct:.1f}%)") - - for i, trace in enumerate(adaptive_traces): - log.info( - f" Trace {i+1}: min_conf={trace['min_conf']:.3f}, " - f"stopped_early={trace.get('stopped_early', False)}, " - f"answer={trace['extracted_answer']}" - ) - - # Combine all traces - all_traces = warmup_traces + adaptive_traces - - # Filter and vote - result = self._filter_and_vote(all_traces, conf_threshold) - - # Build enhanced metadata (consistent with offline mode) - filtered_set = set(id(t) for t in result["filtered_traces"]) - - # Full details for ALL traces (for analysis) - all_traces_details = [] - for i, trace in enumerate(all_traces): - all_traces_details.append( - { - "trace_id": i, - "text": trace["text"], # Full reasoning text - "min_conf": trace["min_conf"], - "mean_conf": ( - sum(trace["token_confs"]) / len(trace["token_confs"]) - if trace["token_confs"] - else 0.0 - ), - "answer": trace["extracted_answer"], - "num_tokens": len(trace.get("token_data", [])), - "selected": id(trace) in filtered_set, - "phase": "warmup" if i < len(warmup_traces) else "adaptive", - "stopped_early": trace.get("stopped_early", False), - "window_confs": trace.get( - "window_confs", [] - ), # For confidence charts - "token_confs": trace.get( - "token_confs", [] - ), # Per-token confidences - } - ) - - # Summary of filtered traces (for quick reference) - filtered_trace_ids = [ - all_traces.index(trace) for trace in result["filtered_traces"] - ] - - # Build metadata using StrategyMetadataBuilder - builder = StrategyMetadataBuilder("deepconf") - - # Add configuration - builder.add_config( - mode="online", - window_size=self.window_size, - temperature=self.temperature, - top_p=self.top_p, - max_tokens=self.max_tokens, - top_logprobs=self.top_logprobs, - filter_method=self.filter_method, - n_threads=self.n_threads, - ) - - # Add results - builder.add_results( - selected_answer=result["selected_answer"], - confidence_score=result["confidence_score"], - vote_distribution=result["vote_distribution"], - total_traces=len(all_traces), - filtered_traces=len(result["filtered_traces"]), - ) - - # Add generation details - builder.add_generation_details( - all_traces=all_traces_details, - filtered_trace_ids=filtered_trace_ids, - ) - - # Add online-specific metadata - builder.add_strategy_specific( - warmup_traces=len(warmup_traces), - adaptive_traces=len(adaptive_traces), - conf_threshold=conf_threshold, - confidence_percentile=self.confidence_percentile, - ) - - # Log summary to console - builder.log_summary(log) - - # Return all traces for analysis, with selected answer for evaluation - return { - "trajectory": result["selected_text"], # Winning trace text - "extracted_answer": result["selected_answer"], # For ExactMatch evaluation - "all_traces": all_traces_details, # All generated traces with full text - "steps": [t["text"] for t in all_traces_details], # All trace texts - "validity_scores": [ - t["min_conf"] for t in all_traces_details - ], # All confidences - "consensus_score": result[ - "agreement_rate" - ], # Proportion of traces with winning answer (0-1) - "vote_distribution": result["vote_distribution"], # Answer -> percentage - "completed": True, - "metadata": builder.build(), - } - - def _generate_single_trace(self, args: tuple) -> Optional[Dict[str, Any]]: - """ - Generate a single trace with logprobs (supports both API and local models). + For each candidate that stopped at , generates the answer phase + via generate_answer_candidates_batch, producing a proper two-step trajectory. Args: - args: Tuple of (prompt, trace_index, total_traces) + request: Chat messages for the request + candidates: List of StepCandidate objects from generation Returns: - Trace dictionary with text, confidence, and answer, or None if error + List of path dictionaries with text, num_tokens, and steps info """ - prompt, i, n = args - try: - # Check if this is an API model (BlackboxModel) or local HuggingFace model - if isinstance(self.model, BlackboxModel): - # ===== API MODEL PATH (OpenAI, OpenRouter, etc.) ===== - # Prepare chat messages - messages = [ - { - "role": "system", - "content": ( - "You are a math problem solver. Always put your " - "final numerical answer in \\boxed{answer} format." - ), - }, - {"role": "user", "content": prompt}, - ] - - # Generate with logprobs - # Disable early stopping for batch generation - self.model.early_stopping = None - - results = self.model.generate_texts( - chats=[messages], - max_new_tokens=self.max_tokens, - temperature=self.temperature, - top_p=self.top_p, - output_scores=True, - top_logprobs=self.top_logprobs, - ) - - # Extract text and logprobs from result - text = results[0]["text"] - logprobs_data = results[0].get("logprobs", []) - - # Extract confidences from logprobs - token_confs = [] - for token_info in logprobs_data: - # Use mean of top-k logprobs (same as online mode) - top_logprobs_list = token_info["top_logprobs"][: self.top_logprobs] - mean_logprob = sum(t["logprob"] for t in top_logprobs_list) / len( - top_logprobs_list - ) - token_confs.append(-mean_logprob) - - token_data = logprobs_data - - else: - # ===== LOCAL HUGGINGFACE MODEL PATH ===== - import torch - import torch.nn.functional as F - - # Get device from model parameters - device = next(self.model.model.parameters()).device - - # Tokenize prompt - inputs = self.model.tokenize([prompt]) - input_ids = inputs["input_ids"].to(device) - attention_mask = inputs["attention_mask"].to(device) - - # Generate with output_scores to get logits - # Use model's generation config for eos_token_id (may be list of tokens) - eos_token_id = self.model.model.generation_config.eos_token_id - if eos_token_id is None: - eos_token_id = self.model.tokenizer.eos_token_id - - with torch.no_grad(): - outputs = self.model.model.generate( - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=self.max_tokens, - temperature=self.temperature, - top_p=self.top_p, - do_sample=True, - return_dict_in_generate=True, - output_scores=True, # Get logits for each generated token - pad_token_id=self.model.tokenizer.pad_token_id, - eos_token_id=eos_token_id, - repetition_penalty=1.1, - ) - - # Extract generated text - output_seq = outputs.sequences[0] - new_tokens = output_seq[input_ids.shape[1] :] - text = self.model.tokenizer.decode(new_tokens, skip_special_tokens=True) - - # Extract logprobs from scores (logits) - logprobs_data = [] - token_confs = [] - - # Check if scores were returned - if outputs.scores is None or len(outputs.scores) == 0: - log.warning( - f" No scores returned by model (scores: {outputs.scores})" - ) - token_data = [] - else: - log.info( - f" Processing {len(outputs.scores)} score tensors for {len(new_tokens)} tokens" - ) - - # Handle length mismatch (scores may not include EOS token) - num_scores = len(outputs.scores) - num_tokens = len(new_tokens) - if num_scores != num_tokens: - log.warning( - f" Length mismatch: {num_scores} scores vs {num_tokens} tokens. Using min length." - ) - process_length = min(num_scores, num_tokens) - else: - process_length = num_scores - - for idx in range(process_length): - score_tensor = outputs.scores[idx] - # score_tensor shape: [batch_size, vocab_size] - # Convert logits to log probabilities - log_probs = F.log_softmax(score_tensor[0], dim=-1) - - # Get top-k log probs and indices - top_k_logprobs, top_k_indices = torch.topk( - log_probs, min(self.top_logprobs, log_probs.size(-1)) - ) - - # Get the token that was actually generated - generated_token_id = new_tokens[idx].item() - generated_logprob = log_probs[generated_token_id].item() - - # Build logprobs data structure (matching API format) - top_logprobs_list = [ - { - "token": self.model.tokenizer.decode([token_id.item()]), - "logprob": logprob.item(), - } - for logprob, token_id in zip(top_k_logprobs, top_k_indices) - ] - - logprobs_data.append( - { - "token": self.model.tokenizer.decode( - [generated_token_id] - ), - "logprob": generated_logprob, - "top_logprobs": top_logprobs_list, - } - ) - - # Compute confidence (negative mean of top-k logprobs) - # Filter out -inf values from masked/impossible tokens - valid_logprobs = top_k_logprobs[~torch.isinf(top_k_logprobs)] - if len(valid_logprobs) > 0: - mean_logprob = valid_logprobs.mean().item() - token_confs.append(-mean_logprob) - else: - # All logprobs are -inf (shouldn't happen, but handle it) - token_confs.append(float("inf")) - - token_data = logprobs_data - - # Clear GPU cache for local models - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - # ===== COMMON PATH (both API and local models) ===== - # Compute sliding window confidences - window_confs = compute_sliding_window_confidence( - token_confs, self.window_size - ) - min_conf = min(window_confs) if window_confs else 0.0 - - # Extract answer - extracted_answer = extract_answer(text) - - trace = { - "text": text, - "min_conf": min_conf, - "extracted_answer": extracted_answer, - "token_confs": token_confs, - "window_confs": window_confs, - "token_data": token_data, - } - + # Identify which candidates need answer generation + thinking_indices = [] + for i, candidate in enumerate(candidates): + if ( + getattr(self.step_generator, "thinking_mode", False) + and candidate.is_thinking_complete + and not candidate.is_trajectory_complete + ): + thinking_indices.append(i) + + # Batch generate all answer phases in one call + answer_map = {} + if thinking_indices: log.info( - f" Trace {i+1}/{n}: tokens={len(token_data)}, " - f"min_conf={min_conf:.3f}, answer={extracted_answer}" + f"Generating {len(thinking_indices)} answer phases in batched call" ) - - # Debug: log text if no answer extracted - if not extracted_answer: - log.warning(f" No answer extracted. Text length: {len(text)}") - log.warning(f" Text ends with: ...{text[-150:]}") - - return trace - - except Exception as e: - log.error(f" Error generating trace {i+1}: {e}") - return None - - def _generate_traces_batch(self, prompt: str, n: int) -> List[Dict[str, Any]]: - """ - Generate n traces using parallel API requests or batched local generation. - - Returns: - List of trace dictionaries with text, confidence, and answer - """ - # Check if this is a vLLM model (fast path) - if hasattr(self.model, "is_vllm") and self.model.is_vllm: - return self._generate_traces_vllm(prompt, n) - - # Check if this is a local HuggingFace model - if not isinstance(self.model, BlackboxModel): - # Use batched generation for local models (much faster!) - return self._generate_traces_batch_local(prompt, n) - - # For API models, use parallel generation - # Prepare arguments for each trace: (prompt, index, total) - trace_args = [(prompt, i, n) for i in range(n)] - - # Use base class parallel generation with our trace-specific worker - return self._parallel_generate( - worker_func=self._generate_single_trace, - task_args=trace_args, - n_threads=self.n_threads, - desc=f"Generating {n} traces", - ) - - def _generate_traces_vllm(self, prompt: str, n: int) -> List[Dict[str, Any]]: - """ - Generate n traces using vLLM (fast, batched, no OOM). - - vLLM generates all traces in a single batched call with: - - PagedAttention for memory efficiency - - Continuous batching for throughput - - Prefix caching for shared prompts - - Returns: - List of trace dictionaries with text, confidence, and answer - """ - from vllm import SamplingParams - - log.info(f"🚀 vLLM batch generation: {n} traces...") - - # Get the vLLM engine - llm = self.model.vllm_engine - tokenizer = llm.get_tokenizer() - - # Prepare the prompt with chat template - messages = [{"role": "user", "content": prompt}] - - # Apply chat template with enable_thinking parameter for Qwen3 models - # Pass enable_thinking=False if disable_thinking_mode is True - formatted_prompt = tokenizer.apply_chat_template( - messages, - tokenize=False, - add_generation_prompt=True, - enable_thinking=not self.disable_thinking_mode, - ) - - # Create sampling params for batch generation - # Use stop strings for early stopping (saves tokens) - sampling_params = SamplingParams( - n=n, # Generate n traces in ONE call - temperature=self.temperature, - top_p=self.top_p, - max_tokens=self.max_tokens, - logprobs=self.top_logprobs, - stop=["", ""], # Early stopping - include_stop_str_in_output=True, # Keep the stop string in output - seed=self.seed, # Reproducibility - ) - - # Single call generates all traces in parallel - log.info( - f" Generating with params: temp={self.temperature}, top_p={self.top_p}, max_tokens={self.max_tokens}" - ) - outputs = llm.generate([formatted_prompt], sampling_params) - - # Process outputs - all_traces = [] - for output in outputs[0].outputs: - text = output.text - - # Extract logprobs and compute confidences - token_confs = [] - logprobs_data = [] - - if output.logprobs: - for i, token_logprobs in enumerate(output.logprobs): - if token_logprobs is None: - continue - - # Get top logprobs for this token - top_logprobs_list = [] - for token_id, logprob_obj in token_logprobs.items(): - top_logprobs_list.append( - { - "token": tokenizer.decode([token_id]), - "logprob": logprob_obj.logprob, - } - ) - - # Sort by logprob (descending) - top_logprobs_list.sort(key=lambda x: x["logprob"], reverse=True) - top_logprobs_list = top_logprobs_list[: self.top_logprobs] - - # Compute confidence (negative mean of top-k logprobs) - if top_logprobs_list: - valid_logprobs = [ - lp["logprob"] - for lp in top_logprobs_list - if lp["logprob"] > float("-inf") - ] - if valid_logprobs: - mean_logprob = sum(valid_logprobs) / len(valid_logprobs) - token_confs.append(-mean_logprob) - - # Store token data - generated_token_id = ( - output.token_ids[i] if i < len(output.token_ids) else None - ) - logprobs_data.append( - { - "token": ( - tokenizer.decode([generated_token_id]) - if generated_token_id - else "" - ), - "logprob": ( - top_logprobs_list[0]["logprob"] - if top_logprobs_list - else 0 - ), - "top_logprobs": top_logprobs_list, - } - ) - - # Compute sliding window confidences - window_confs = compute_sliding_window_confidence( - token_confs, self.window_size + batch_requests = [request] * len(thinking_indices) + batch_trajectories = [[candidates[i]] for i in thinking_indices] + answer_results = self.step_generator.generate_answer_candidates_batch( + batch_requests, + batch_trajectories, + candidates_per_step=1, ) - min_conf = min(window_confs) if window_confs else 0.0 - - # Extract answer - extracted_answer = extract_answer(text) - - trace = { - "text": text, - "min_conf": min_conf, - "extracted_answer": extracted_answer, - "token_confs": token_confs, - "window_confs": window_confs, - "token_data": logprobs_data, - } - - all_traces.append(trace) - - log.info( - f" Trace {len(all_traces)}/{n}: tokens={len(logprobs_data)}, " - f"min_conf={min_conf:.3f}, answer={extracted_answer}" + for batch_idx, cand_idx in enumerate(thinking_indices): + if answer_results[batch_idx]: + answer_map[cand_idx] = answer_results[batch_idx][0] + + # Build paths from candidates + answers + paths = [] + for i, candidate in enumerate(candidates): + text = candidate.raw_text or candidate.text + num_tokens = (candidate.other_data or {}).get( + "original_token_count", len(candidate.token_ids) ) - # Summary statistics - total_tokens = sum(len(t.get("token_data", [])) for t in all_traces) - avg_tokens = total_tokens / len(all_traces) if all_traces else 0 - - log.info(f"✓ vLLM generation complete: {len(all_traces)}/{n} traces") - log.info( - f" Total tokens: {total_tokens}, Average: {avg_tokens:.0f} tokens/trace" - ) - - return all_traces - - def _generate_traces_batch_local(self, prompt: str, n: int) -> List[Dict[str, Any]]: - """ - Generate n traces using SEQUENTIAL generation with early stopping. - - Uses AnswerStoppingCriteria to stop generation as soon as - a complete answer is found ( for default format, - or \\boxed{...} for boxed format), saving significant tokens. - - Sequential generation is used because HuggingFace's StoppingCriteria - with batched generation (num_return_sequences > 1) only stops when - ALL sequences are done, negating any early stopping benefits. - - Returns: - List of trace dictionaries with text, confidence, and answer - """ - import torch - import torch.nn.functional as F - - # Detect answer format from prompt - answer_format = ( - "default" - if ":" in prompt or "" in prompt - else "boxed" - ) - log.info(f"🚀 Sequential generation with early stopping: {n} sequences...") - - all_traces = [] - - # Get device from model parameters - device = next(self.model.model.parameters()).device - - # Tokenize prompt once for all sequences - inputs = self.model.tokenize([prompt]) - input_ids = inputs["input_ids"].to(device) - attention_mask = inputs["attention_mask"].to(device) - start_length = input_ids.shape[1] - - # Get eos_token_id from model config - eos_token_id = self.model.model.generation_config.eos_token_id - if eos_token_id is None: - eos_token_id = self.model.tokenizer.eos_token_id - - # Generate each sequence independently (allows per-sequence early stopping) - for seq_idx in range(n): - try: - # Create fresh stopping criteria for each sequence - stopping_criteria = AnswerStoppingCriteria( - tokenizer=self.model.tokenizer, - start_length=start_length, - batch_size=1, # Single sequence - answer_format=answer_format, - min_tokens_after=0, # Stop immediately when answer found - ) - - with torch.no_grad(): - outputs = self.model.model.generate( - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=self.max_tokens, - temperature=self.temperature, - top_p=self.top_p, - do_sample=True, - num_return_sequences=1, # Single sequence for early stopping - return_dict_in_generate=True, - output_scores=True, # Get logits for each generated token - pad_token_id=self.model.tokenizer.pad_token_id, - eos_token_id=eos_token_id, - repetition_penalty=1.1, - stopping_criteria=StoppingCriteriaList([stopping_criteria]), - ) - - # Extract generated text - output_seq = outputs.sequences[0] - new_tokens = output_seq[start_length:] - text = self.model.tokenizer.decode(new_tokens, skip_special_tokens=True) - - # Check if we stopped early - stopped_early = stopping_criteria.finished[0] - - # Extract logprobs from scores - logprobs_data = [] - token_confs = [] - - if outputs.scores is None or len(outputs.scores) == 0: - log.warning(f" Sequence {seq_idx+1}/{n}: No scores returned") - token_data = [] - else: - num_scores = len(outputs.scores) - num_tokens = len(new_tokens) - process_length = min(num_scores, num_tokens) - - for token_idx in range(process_length): - score_tensor = outputs.scores[token_idx] - # score_tensor shape: [1, vocab_size] (single sequence) - log_probs = F.log_softmax(score_tensor[0], dim=-1) - - # Get top-k log probs and indices - top_k_logprobs, top_k_indices = torch.topk( - log_probs, min(self.top_logprobs, log_probs.size(-1)) - ) - - # Get the token that was actually generated - generated_token_id = new_tokens[token_idx].item() - generated_logprob = log_probs[generated_token_id].item() - - # Build logprobs data structure (matching API format) - top_logprobs_list = [ - { - "token": self.model.tokenizer.decode([token_id.item()]), - "logprob": logprob.item(), - } - for logprob, token_id in zip(top_k_logprobs, top_k_indices) - ] - - logprobs_data.append( - { - "token": self.model.tokenizer.decode( - [generated_token_id] - ), - "logprob": generated_logprob, - "top_logprobs": top_logprobs_list, - } - ) - - # Compute confidence (negative mean of top-k logprobs) - valid_logprobs = top_k_logprobs[~torch.isinf(top_k_logprobs)] - if len(valid_logprobs) > 0: - mean_logprob = valid_logprobs.mean().item() - token_confs.append(-mean_logprob) - else: - token_confs.append(float("inf")) - - token_data = logprobs_data - - # Compute sliding window confidences - window_confs = compute_sliding_window_confidence( - token_confs, self.window_size - ) - min_conf = min(window_confs) if window_confs else 0.0 - - # Extract answer - extracted_answer = extract_answer(text) - - trace = { - "text": text, - "min_conf": min_conf, - "extracted_answer": extracted_answer, - "token_confs": token_confs, - "window_confs": window_confs, - "token_data": token_data, - "stopped_early": stopped_early, - } - - early_str = " [EARLY STOP]" if stopped_early else "" - log.info( - f" Trace {seq_idx+1}/{n}: tokens={len(token_data)}, " - f"min_conf={min_conf:.3f}, answer={extracted_answer}{early_str}" - ) - log.info( - f" --- Trace {seq_idx+1} text ---\n{text}\n --- End trace {seq_idx+1} ---" + if i in answer_map: + answer_step = answer_map[i] + answer_step.is_trajectory_complete = True + trajectory = [candidate, answer_step] + full_text = convert_trajectory_to_string(trajectory) + answer_tokens = ( + len(answer_step.token_ids) if answer_step.token_ids else 0 ) - - all_traces.append(trace) - - except Exception as e: - log.error(f" Error generating sequence {seq_idx+1}/{n}: {e}") - import traceback - - log.error(traceback.format_exc()) - - # Clear GPU cache periodically - if torch.cuda.is_available() and (seq_idx + 1) % 4 == 0: - torch.cuda.empty_cache() - - # Summary statistics - num_early_stopped = sum(1 for t in all_traces if t.get("stopped_early", False)) - total_tokens = sum(len(t.get("token_data", [])) for t in all_traces) - avg_tokens = total_tokens / len(all_traces) if all_traces else 0 - - log.info(f"✓ Sequential generation complete: {len(all_traces)}/{n} sequences") - log.info(f" Early stopped: {num_early_stopped}/{len(all_traces)}") - log.info(f" Average tokens: {avg_tokens:.0f} (max: {self.max_tokens})") - if num_early_stopped > 0: - savings = (1 - avg_tokens / self.max_tokens) * 100 - log.info(f" Token savings: {savings:.1f}%") - - return all_traces - - def _generate_traces_adaptive( - self, prompt: str, n: int, conf_threshold: float - ) -> List[Dict[str, Any]]: - """ - Generate n traces with adaptive early stopping based on confidence. - - Uses streaming + logprobs to stop generation when confidence drops. - """ - log.info( - f"Generating {n} adaptive traces with confidence " - f"threshold={conf_threshold:.3f}" - ) - - traces = [] - - for i in range(n): - try: - # Prepare messages - messages = [ + num_tokens += answer_tokens + paths.append( { - "role": "system", - "content": ( - "You are a math problem solver. Always put your " - "final numerical answer in \\boxed{answer} format." - ), - }, - {"role": "user", "content": prompt}, - ] - - # Configure model with confidence-based early stopping - self.model.early_stopping = ConfidenceEarlyStopping( - threshold=conf_threshold, - window_size=self.window_size, - top_k=self.top_logprobs, - method="mean_logprob", # Same formula as warmup phase - ) - - # Generate with early stopping configured in model - results = self.model.generate_texts( - [messages], - max_new_tokens=self.max_tokens, - temperature=self.temperature, - ) - - result = results[0] - text = result["text"] - logprobs_data = result["logprobs"] - stopped_early = result.get("stopped_early", False) - - # Compute token confidences from logprobs - token_confs = [] - for token_info in logprobs_data: - conf = compute_token_confidence_from_logprobs( - token_info["top_logprobs"], topk=self.top_logprobs - ) - token_confs.append(conf) - - # Compute sliding window confidences - window_confs = compute_sliding_window_confidence( - token_confs, self.window_size + "text": full_text, + "num_tokens": num_tokens, + "steps": [candidate.text, answer_step.text], + "is_complete": True, + } ) - min_conf = min(window_confs) if window_confs else 0.0 - - # Extract answer - extracted_answer = extract_answer(text) + continue - trace = { + # Non-thinking or no : single-step path + paths.append( + { "text": text, - "min_conf": min_conf, - "extracted_answer": extracted_answer, - "token_confs": token_confs, - "window_confs": window_confs, - "token_data": logprobs_data, - "stopped_early": stopped_early, + "num_tokens": num_tokens, + "steps": [text], + "is_complete": candidate.is_trajectory_complete, } + ) - traces.append(trace) - - early_str = " (stopped early)" if stopped_early else "" - log.info( - f" Trace {i+1}/{n}: tokens={len(logprobs_data)}, " - f"min_conf={min_conf:.3f}, answer={extracted_answer}{early_str}" - ) - - # Debug: log text if no answer extracted - if not extracted_answer: - log.warning(f" No answer extracted. Text length: {len(text)}") - log.warning(f" Text ends with: ...{text[-150:]}") - - except Exception as e: - log.error(f" Error generating adaptive trace {i+1}: {e}") - import traceback + return paths - log.error(traceback.format_exc()) - continue + def _compute_trace_confidence(self, raw_logprobs, window_size, top_k=20): + """Compute sliding-window min confidence from vLLM raw_logprobs. - return traces - - def _extract_logprobs_from_model(self) -> tuple: - """ - Extract logprobs from model's stored data. + Args: + raw_logprobs: List[Dict[token_id -> Logprob]] from vLLM. + Each dict maps token IDs to Logprob objects with .logprob attribute. + window_size: Sliding window size. + top_k: Number of top logprobs to use for confidence. Returns: - (token_confidences, token_data) + (min_conf, token_confs, window_confs) tuple. """ token_confs = [] - token_data = [] - - if hasattr(self.model, "logprobs") and len(self.model.logprobs) > 0: - logprobs_obj = self.model.logprobs[0] - - if hasattr(logprobs_obj, "content") and logprobs_obj.content: - for token_info in logprobs_obj.content: - # Build token data - token_entry = { - "token": token_info.token, - "logprob": token_info.logprob, - "top_logprobs": [ - {"token": t.token, "logprob": t.logprob} - for t in token_info.top_logprobs - ], - } - token_data.append(token_entry) - - # Compute confidence - conf = compute_token_confidence_from_logprobs( - token_entry["top_logprobs"], topk=self.top_logprobs - ) - token_confs.append(conf) + for token_logprobs_dict in raw_logprobs: + if token_logprobs_dict is None: + token_confs.append(float("inf")) + continue + # token_logprobs_dict: Dict[token_id -> Logprob] + logprob_values = sorted( + [lp.logprob for lp in token_logprobs_dict.values()], + reverse=True, + )[:top_k] + valid = [lp for lp in logprob_values if lp > float("-inf")] + if valid: + token_confs.append(-sum(valid) / len(valid)) + else: + token_confs.append(float("inf")) - return token_confs, token_data + window_confs = compute_sliding_window_confidence(token_confs, window_size) + min_conf = min(window_confs) if window_confs else 0.0 + return min_conf, token_confs, window_confs def _filter_and_vote( self, traces: List[Dict[str, Any]], conf_threshold: Optional[float] = None @@ -1228,11 +218,11 @@ def _filter_and_vote( Filter traces by confidence and perform majority voting. Args: - traces: List of trace dictionaries - conf_threshold: Confidence threshold (optional) + traces: List of trace dictionaries with min_conf, extracted_answer, text. + conf_threshold: Confidence threshold (optional). Returns: - Dictionary with selected answer and metadata + Dictionary with selected answer and metadata. """ if not traces: return { @@ -1241,28 +231,26 @@ def _filter_and_vote( "confidence_score": 0.0, "filtered_traces": [], "vote_distribution": {}, + "agreement_rate": 0.0, } # Apply filtering if self.filter_method.startswith("top"): - # TopK filtering - k = int(self.filter_method[3:]) # e.g., "top10" -> 10 + k = int(self.filter_method[3:]) filtered = sorted(traces, key=lambda t: t["min_conf"], reverse=True)[:k] - log.info(f"🔍 Top-{k} filtering: {len(filtered)}/{len(traces)} traces") + log.info(f"Top-{k} filtering: {len(filtered)}/{len(traces)} traces") elif conf_threshold is not None or self.confidence_threshold is not None: - # Threshold filtering threshold = conf_threshold or self.confidence_threshold filtered = [t for t in traces if t["min_conf"] >= threshold] log.info( - f"🔍 Threshold filtering (≥{threshold:.3f}): " + f"Threshold filtering (>={threshold:.3f}): " f"{len(filtered)}/{len(traces)} traces" ) else: - # No filtering filtered = traces - log.info(f"🔍 No filtering: {len(filtered)} traces") + log.info(f"No filtering: {len(filtered)} traces") if not filtered: log.warning("No traces passed filter, using all traces") @@ -1280,6 +268,7 @@ def _filter_and_vote( "confidence_score": 0.0, "filtered_traces": filtered, "vote_distribution": {}, + "agreement_rate": 0.0, } # Sum weights per answer @@ -1291,8 +280,7 @@ def _filter_and_vote( selected_answer = max(answer_weights.keys(), key=lambda x: answer_weights[x]) total_weight = sum(answer_weights.values()) - # Confidence = mean of all trace confidences (not voting-based) - # Based on deepconf/utils.py:219 + # Confidence = mean of all trace confidences confidence_score = sum(weights) / len(weights) if weights else 0.0 # Get trace with selected answer @@ -1307,8 +295,7 @@ def _filter_and_vote( ans: weight / total_weight for ans, weight in answer_weights.items() } - # Calculate answer statistics - num_unique_answers = len(answer_weights) + # Agreement rate total_answers = len(answers) agreement_rate = ( max(answers.count(ans) for ans in answer_weights.keys()) / total_answers @@ -1316,24 +303,385 @@ def _filter_and_vote( else 0 ) - log.info(f"🏆 Selected answer: '{selected_answer}'") + log.info(f"Selected answer: '{selected_answer}'") log.info(f" Confidence: {confidence_score:.3f}") - log.info(f" Answers: {total_answers} total, {num_unique_answers} unique") - log.info( - f" Agreement: {agreement_rate:.1%} \ - ({max(answers.count(ans) for ans in answer_weights.keys())}/{total_answers} traces)" - ) - log.info(" Vote distribution (weighted by min_conf):") + log.info(f" Agreement: {agreement_rate:.1%}") for ans, pct in sorted(vote_dist.items(), key=lambda x: x[1], reverse=True): - raw_weight = answer_weights[ans] count = answers.count(ans) - log.info(f" {ans}: {pct:.1%} (weight={raw_weight:.3f}, count={count})") + log.info(f" {ans}: {pct:.1%} (count={count})") return { "selected_answer": selected_answer, "selected_text": selected_trace["text"] if selected_trace else "", "confidence_score": confidence_score, - "agreement_rate": agreement_rate, # Proportion of traces with winning answer (0-1) + "agreement_rate": agreement_rate, "filtered_traces": filtered, "vote_distribution": vote_dist, } + + def _process_candidates( + self, + request: List[Dict[str, str]], + candidates: List[Any], + ) -> List[Dict[str, Any]]: + """ + Process candidates into trace dicts with confidence and answers. + + Handles thinking path completion, confidence computation, and answer + extraction. Used by both offline and online modes. + + Args: + request: Chat messages for the request. + candidates: List of StepCandidate objects. + + Returns: + List of trace dicts with text, min_conf, extracted_answer, + token_confs, window_confs, num_tokens. + """ + paths = self._complete_thinking_paths(request, candidates) + + traces = [] + for candidate, path in zip(candidates, paths): + raw_logprobs = (candidate.other_data or {}).get("raw_logprobs", []) + min_conf, token_confs, window_confs = self._compute_trace_confidence( + raw_logprobs, self.window_size + ) + + extracted_answer = extract_answer(path["text"], answer_format="auto") + + traces.append( + { + "text": path["text"], + "min_conf": min_conf, + "extracted_answer": extracted_answer, + "token_confs": token_confs, + "window_confs": window_confs, + "num_tokens": path["num_tokens"], + } + ) + + return traces + + def _build_result( + self, + traces: List[Dict[str, Any]], + sample_idx: int, + token_stats: Dict[str, Any], + mode: str = "offline", + early_stopped: bool = False, + ) -> Dict[str, Any]: + """ + Build result dict from processed traces. + + Logs per-trace info, runs filter/vote, and assembles the final result + dictionary. Used by both offline and online modes. + + Args: + traces: List of trace dicts from _process_candidates. + sample_idx: Sample index for logging. + token_stats: Token statistics dict. + mode: "offline" or "online". + early_stopped: Whether the sample stopped early (online mode). + + Returns: + Result dictionary with trajectory, metadata, etc. + """ + if not traces: + return self._empty_result() + + # Log trace info + for i, trace in enumerate(traces): + log.info( + f" Sample {sample_idx}, Trace {i + 1}/{len(traces)}: " + f"tokens={trace['num_tokens']}, min_conf={trace['min_conf']:.3f}, " + f"answer={trace['extracted_answer']}" + ) + + # Filter and vote + vote_result = self._filter_and_vote(traces) + + token_stats["generation_count"] = len(traces) + + # Build metadata + filtered_set = set(id(t) for t in vote_result["filtered_traces"]) + + all_traces_details = [] + for i, trace in enumerate(traces): + all_traces_details.append( + { + "trace_id": i, + "text": trace["text"], + "min_conf": trace["min_conf"], + "mean_conf": ( + sum(trace["token_confs"]) / len(trace["token_confs"]) + if trace["token_confs"] + else 0.0 + ), + "answer": trace["extracted_answer"], + "num_tokens": trace["num_tokens"], + "selected": id(trace) in filtered_set, + "window_confs": trace.get("window_confs", []), + "token_confs": trace.get("token_confs", []), + } + ) + + builder = StrategyMetadataBuilder("deepconf") + config_kwargs = dict( + budget=self.budget, + window_size=self.window_size, + filter_method=self.filter_method, + mode=mode, + ) + if mode == "online": + config_kwargs["online_batch_size"] = self.online_batch_size + config_kwargs["min_agreement"] = self.min_agreement + builder.add_config(**config_kwargs) + builder.add_results( + selected_answer=vote_result["selected_answer"], + confidence_score=vote_result["confidence_score"], + vote_distribution=vote_result["vote_distribution"], + total_traces=len(traces), + filtered_traces=len(vote_result["filtered_traces"]), + ) + builder.add_generation_details( + all_traces=all_traces_details, + ) + + result = { + "trajectory": vote_result["selected_text"], + "steps": [t["text"] for t in all_traces_details], + "validity_scores": [t["min_conf"] for t in all_traces_details], + "completed": bool(traces), + "strategy": "deepconf", + "extracted_answer": vote_result["selected_answer"], + "consensus_score": vote_result["agreement_rate"], + "vote_distribution": vote_result["vote_distribution"], + "metadata": builder.build(), + "all_traces": all_traces_details, + "total_tokens": sum(t["num_tokens"] for t in traces), + "token_stats": token_stats, + } + + if mode == "online": + result["early_stopped"] = early_stopped + + return result + + def generate_trajectories_batch( + self, + requests: List[List[Dict[str, str]]], + sample_indices: List[int] = None, + save_callback: Callable = None, + ) -> List[Dict[str, Any]]: + """ + Generate N traces for each of M samples using step_generator. + + Dispatches to offline or online mode based on self.mode. + + Args: + requests: List of M chat message lists (one per sample). + sample_indices: Optional list of sample indices for logging. + save_callback: Optional callback for progressive saves. + + Returns: + List of M result dictionaries (one per sample). + """ + if sample_indices is None: + sample_indices = list(range(len(requests))) + + if self.mode == "online": + return self._generate_online(requests, sample_indices, save_callback) + else: + return self._generate_offline(requests, sample_indices, save_callback) + + def _build_stop_tokens(self): + """Build stop tokens list based on generator config.""" + stop_tokens = [""] + if ( + getattr(self.step_generator, "thinking_mode", False) + and "" not in stop_tokens + ): + stop_tokens.append("") + return stop_tokens + + def _generate_offline( + self, + requests: List[List[Dict[str, str]]], + sample_indices: List[int], + save_callback: Callable = None, + ) -> List[Dict[str, Any]]: + """ + Generate traces in offline mode (single batched call). + + All M x N trajectories are generated in one call, then confidence is + computed from raw_logprobs and filtering + voting is applied per sample. + """ + M = len(requests) + N = self.budget + + log.info( + f"DeepConf offline: generating {M} samples x {N} traces = {M * N} " + f"trajectories via generate_step_candidates_batch" + ) + + stop_tokens = self._build_stop_tokens() + + # Single batched call: M samples x N traces + self.step_generator.reset_per_sample_stats() + batch_results = self.step_generator.generate_step_candidates_batch( + requests=requests, + trajectories=[[]] * M, + candidates_per_step=N, + stop_tokens_override=stop_tokens, + max_tokens=self.step_generator.generation_limit, + compute_uncertainty=False, + sample_ids=list(range(M)), + ) + + results = [] + for idx, (candidates, sample_idx) in enumerate( + zip(batch_results, sample_indices) + ): + if not candidates: + log.error(f"No output generated for sample {sample_idx}") + results.append(self._empty_result()) + continue + + traces = self._process_candidates(requests[idx], candidates) + token_stats = self.step_generator.get_sample_stats_for(idx) + result = self._build_result(traces, sample_idx, token_stats) + results.append(result) + + log.info( + f"DeepConf offline: completed {len(results)} samples, " + f"total {M * N} trajectories generated" + ) + return results + + def _generate_online( + self, + requests: List[List[Dict[str, str]]], + sample_indices: List[int], + save_callback: Callable = None, + ) -> List[Dict[str, Any]]: + """ + Generate traces in online mode with early stopping. + + Generates traces in small batches per round. After each round, + checks agreement among filtered traces for each sample. Samples + that reach sufficient agreement or exhaust their budget stop + participating in future rounds. + """ + M = len(requests) + + log.info( + f"DeepConf online: {M} samples, batch_size={self.online_batch_size}, " + f"budget={self.budget}, min_agreement={self.min_agreement}" + ) + + stop_tokens = self._build_stop_tokens() + + # Per-sample state + all_traces = [[] for _ in range(M)] + completed = [False] * M + early_stopped = [False] * M + + self.step_generator.reset_per_sample_stats() + round_num = 0 + + while not all(completed): + round_num += 1 + active_indices = [i for i in range(M) if not completed[i]] + + log.info( + f"Online round {round_num}: generating for {len(active_indices)} " + f"active samples" + ) + + # Generate batch for active samples only + active_requests = [requests[i] for i in active_indices] + + batch_results = self.step_generator.generate_step_candidates_batch( + requests=active_requests, + trajectories=[[]] * len(active_indices), + candidates_per_step=self.online_batch_size, + stop_tokens_override=stop_tokens, + max_tokens=self.step_generator.generation_limit, + compute_uncertainty=False, + sample_ids=active_indices, + ) + + for batch_idx, sample_idx in enumerate(active_indices): + candidates = batch_results[batch_idx] + if not candidates: + log.warning( + f"No candidates for sample {sample_indices[sample_idx]} " + f"in round {round_num}" + ) + continue + + # Process new candidates and add to pool + new_traces = self._process_candidates( + requests[sample_idx], candidates + ) + all_traces[sample_idx].extend(new_traces) + current_count = len(all_traces[sample_idx]) + + # Check stopping conditions + if current_count >= self.budget: + # Trim to budget if overshot + all_traces[sample_idx] = all_traces[sample_idx][: self.budget] + completed[sample_idx] = True + log.info( + f"Sample {sample_indices[sample_idx]}: budget reached " + f"({self.budget} traces)" + ) + else: + # Check agreement among filtered traces + vote_result = self._filter_and_vote(all_traces[sample_idx]) + agreement = vote_result["agreement_rate"] + if agreement >= self.min_agreement: + completed[sample_idx] = True + early_stopped[sample_idx] = True + log.info( + f"Sample {sample_indices[sample_idx]}: early stop at " + f"{current_count} traces, agreement={agreement:.0%}" + ) + + # Build final results + results = [] + for i in range(M): + token_stats = self.step_generator.get_sample_stats_for(i) + result = self._build_result( + all_traces[i], + sample_indices[i], + token_stats, + mode="online", + early_stopped=early_stopped[i], + ) + results.append(result) + + total_traces = sum(len(t) for t in all_traces) + early_count = sum(early_stopped) + log.info( + f"DeepConf online: completed {M} samples, {total_traces} total traces " + f"({early_count} early stopped)" + ) + return results + + def _empty_result(self) -> Dict[str, Any]: + """Return empty result for failed generation.""" + return { + "trajectory": "", + "steps": [], + "validity_scores": [], + "completed": False, + "strategy": "deepconf", + "extracted_answer": "", + "consensus_score": 0.0, + "vote_distribution": {}, + "metadata": {}, + "all_traces": [], + "total_tokens": 0, + "token_stats": {}, + } diff --git a/scripts/run_tts_eval.py b/scripts/run_tts_eval.py index 7aa69290..23a114a6 100644 --- a/scripts/run_tts_eval.py +++ b/scripts/run_tts_eval.py @@ -450,131 +450,128 @@ def create_model(config): log.info("vLLM model loaded successfully") - # Create step generator for strategies that need it - # DeepConf has its own generation logic - step_generator = None - if config.strategy.type not in ("deepconf",): - if not VLLM_GENERATOR_AVAILABLE: + # Create step generator + if not VLLM_GENERATOR_AVAILABLE: + raise ImportError( + "vLLM step generator not available. " + "Ensure llm_tts.step_candidate_generator_through_vllm is installed." + ) + + # Self-consistency, baseline, deepconf don't need uncertainty wrapper + if config.strategy.type in ( + "self_consistency", + "baseline", + "extended_thinking", + "deepconf", + ): + vllm_model = llm + log.info( + f"{config.strategy.type}: using raw vLLM (no uncertainty wrapper)" + ) + else: + if not POLYGRAPH_UNCERTAINTY_AVAILABLE: raise ImportError( - "vLLM step generator not available. " - "Ensure llm_tts.step_candidate_generator_through_vllm is installed." + "lm-polygraph uncertainty components not available. " + "Ensure lm_polygraph_updates package is installed." ) - # Self-consistency and baseline don't need uncertainty wrapper - # (self-consistency uses majority voting, baseline uses raw vLLM batch generation) - if config.strategy.type in ( - "self_consistency", - "baseline", - "extended_thinking", - ): - vllm_model = llm - log.info( - f"{config.strategy.type}: using raw vLLM (no uncertainty wrapper)" - ) + # Select estimator based on scorer config + scorer_type = config.scorer.type if config.scorer else "entropy" + if scorer_type == "perplexity": + stat_calculators = [VLLMLogprobsCalculator()] + estimator = Perplexity() + elif scorer_type == "sequence_prob": + # Sequence probability scoring (sum of log-probs, not normalized) + stat_calculators = [VLLMLogprobsCalculator()] + estimator = MaximumSequenceProbability() + elif scorer_type == "uncertainty_pd": + # PD-Gap scoring using top-k logprobs matrix + from llm_tts.scorers.estimator_uncertainty_pd import PDGap + + stat_calculators = [VLLMLogprobsCalculator(output_matrix=True)] + estimator = PDGap() + elif scorer_type == "entropy": + # Entropy-based scoring + stat_calculators = [VLLMLogprobsCalculator(), EntropyCalculator()] + estimator = MeanTokenEntropy() + elif scorer_type == "prm": + # PRM scorer uses its own model for scoring + # Use entropy wrapper for generation (scores not used for selection) + stat_calculators = [VLLMLogprobsCalculator(), EntropyCalculator()] + estimator = MeanTokenEntropy() else: - if not POLYGRAPH_UNCERTAINTY_AVAILABLE: - raise ImportError( - "lm-polygraph uncertainty components not available. " - "Ensure lm_polygraph_updates package is installed." - ) - - # Select estimator based on scorer config - scorer_type = config.scorer.type if config.scorer else "entropy" - if scorer_type == "perplexity": - stat_calculators = [VLLMLogprobsCalculator()] - estimator = Perplexity() - elif scorer_type == "sequence_prob": - # Sequence probability scoring (sum of log-probs, not normalized) - stat_calculators = [VLLMLogprobsCalculator()] - estimator = MaximumSequenceProbability() - elif scorer_type == "uncertainty_pd": - # PD-Gap scoring using top-k logprobs matrix - from llm_tts.scorers.estimator_uncertainty_pd import PDGap - - stat_calculators = [VLLMLogprobsCalculator(output_matrix=True)] - estimator = PDGap() - elif scorer_type == "entropy": - # Entropy-based scoring - stat_calculators = [VLLMLogprobsCalculator(), EntropyCalculator()] - estimator = MeanTokenEntropy() - elif scorer_type == "prm": - # PRM scorer uses its own model for scoring - # Use entropy wrapper for generation (scores not used for selection) - stat_calculators = [VLLMLogprobsCalculator(), EntropyCalculator()] - estimator = MeanTokenEntropy() - else: - raise ValueError( - f"Unsupported scorer type for vLLM: {scorer_type}. " - f"Supported types: perplexity, sequence_prob, uncertainty_pd, entropy, prm" - ) - - vllm_model = VLLMWithUncertainty( - llm=llm, - stat_calculators=stat_calculators, - estimator=estimator, - ) - log.info( - f"Created VLLMWithUncertainty wrapper with {type(estimator).__name__}" + raise ValueError( + f"Unsupported scorer type for vLLM: {scorer_type}. " + f"Supported types: perplexity, sequence_prob, uncertainty_pd, entropy, prm" ) - # Always use ThinkingMarkerDetector for step boundary detection - # Stop tokens are derived from detector's semantic markers - # thinking_mode controls two-phase generation (...) - # Logic for disable_thinking_mode: - # None = model doesn't support thinking (e.g., Qwen2.5-Math) -> thinking_mode=False - # False = model supports thinking, enabled (e.g., Qwen3) -> thinking_mode=True - # True = model supports thinking, disabled -> thinking_mode=False - disable_thinking_mode = config.model.get("disable_thinking_mode", None) - thinking_mode = disable_thinking_mode is False + vllm_model = VLLMWithUncertainty( + llm=llm, + stat_calculators=stat_calculators, + estimator=estimator, + ) log.info( - f"Creating VLLMStepGenerator with ThinkingMarkerDetector " - f"(thinking_mode={thinking_mode})" + f"Created VLLMWithUncertainty wrapper with {type(estimator).__name__}" ) - if "min_step_tokens" not in config.strategy: - log.warning("strategy.min_step_tokens not set, defaulting to 50") - if "max_step_tokens" not in config.strategy: - log.warning("strategy.max_step_tokens not set, defaulting to 300") - detector = ThinkingMarkerDetector( - min_step_tokens=config.strategy.get("min_step_tokens", 50), - max_step_tokens=config.strategy.get("max_step_tokens", 300), - use_sequence=config.strategy.get("use_sequence", True), - use_conclusion=config.strategy.get("use_conclusion", True), - use_thinking=config.strategy.get("use_thinking", True), - use_verification=config.strategy.get("use_verification", True), - use_reasoning=config.strategy.get("use_reasoning", False), - use_correction=config.strategy.get("use_correction", False), - use_structure=config.strategy.get("use_structure", False), - custom_markers=config.strategy.get("custom_markers", None), - ) + # Always use ThinkingMarkerDetector for step boundary detection + # Stop tokens are derived from detector's semantic markers + # thinking_mode controls two-phase generation (...) + # Logic for disable_thinking_mode: + # None = model doesn't support thinking (e.g., Qwen2.5-Math) -> thinking_mode=False + # False = model supports thinking, enabled (e.g., Qwen3) -> thinking_mode=True + # True = model supports thinking, disabled -> thinking_mode=False + disable_thinking_mode = config.model.get("disable_thinking_mode", None) + thinking_mode = disable_thinking_mode is False + log.info( + f"Creating VLLMStepGenerator with ThinkingMarkerDetector " + f"(thinking_mode={thinking_mode})" + ) - # Stop token IDs (e.g., [151645, 151643] for Qwen EOS) - # Stop tokens are derived from detector's use_* flags automatically - stop_token_ids = config.strategy.get("stop_token_ids", None) - if stop_token_ids is not None: - stop_token_ids = list(stop_token_ids) - - step_generator = VLLMStepGenerator( - model=vllm_model, - thinking_mode=thinking_mode, - detector=detector, - stop_token_ids=stop_token_ids, - max_new_tokens=config.generation.max_new_tokens, - temperature=config.generation.temperature, - top_p=config.generation.top_p, - top_k=config.generation.get("top_k", 20), - presence_penalty=config.generation.get("presence_penalty", 0.0), - answer_patterns=config.strategy.get( - "detector_answer_patterns", - [], # Empty by default - rely on EOS token IDs - ), - max_context_budget=config.model.get( - "max_context_budget", config.model.get("max_model_len", 32768) - ), - disable_thinking_mode=config.model.get("disable_thinking_mode", None), - ) + if "min_step_tokens" not in config.strategy: + log.warning("strategy.min_step_tokens not set, defaulting to 50") + if "max_step_tokens" not in config.strategy: + log.warning("strategy.max_step_tokens not set, defaulting to 300") + detector = ThinkingMarkerDetector( + min_step_tokens=config.strategy.get("min_step_tokens", 50), + max_step_tokens=config.strategy.get("max_step_tokens", 300), + use_sequence=config.strategy.get("use_sequence", True), + use_conclusion=config.strategy.get("use_conclusion", True), + use_thinking=config.strategy.get("use_thinking", True), + use_verification=config.strategy.get("use_verification", True), + use_reasoning=config.strategy.get("use_reasoning", False), + use_correction=config.strategy.get("use_correction", False), + use_structure=config.strategy.get("use_structure", False), + custom_markers=config.strategy.get("custom_markers", None), + ) + + # Stop token IDs (e.g., [151645, 151643] for Qwen EOS) + # Stop tokens are derived from detector's use_* flags automatically + stop_token_ids = config.strategy.get("stop_token_ids", None) + if stop_token_ids is not None: + stop_token_ids = list(stop_token_ids) + + step_generator = VLLMStepGenerator( + model=vllm_model, + thinking_mode=thinking_mode, + detector=detector, + stop_token_ids=stop_token_ids, + max_new_tokens=config.generation.max_new_tokens, + temperature=config.generation.temperature, + top_p=config.generation.top_p, + top_k=config.generation.get("top_k", 20), + presence_penalty=config.generation.get("presence_penalty", 0.0), + answer_patterns=config.strategy.get( + "detector_answer_patterns", + [], # Empty by default - rely on EOS token IDs + ), + max_context_budget=config.model.get( + "max_context_budget", config.model.get("max_model_len", 32768) + ), + disable_thinking_mode=config.model.get("disable_thinking_mode", None), + ) - log.info(f"Created vLLM step generator: {type(step_generator).__name__}") + log.info(f"Created vLLM step generator: {type(step_generator).__name__}") return model, step_generator @@ -661,128 +658,117 @@ def create_model(config): api_key = config.model.get("api_key") or os.getenv("OPENAI_API_KEY") base_url = config.model.get("base_url", None) - # Check if DeepConf strategy - if config.strategy.type == "deepconf": - # DeepConf uses streaming with logprobs but no boundary detector - model = BlackboxModelWithStreaming( - openai_api_key=api_key, - model_path=model_path, - supports_logprobs=True, - base_url=base_url, - ) - step_generator = None # DeepConf doesn't use step generator - else: - # Other strategies use boundary detection via early stopping - from lm_polygraph.utils import APIWithUncertainty - - from llm_tts.early_stopping import BoundaryEarlyStopping - - # Determine thinking mode (same logic as vLLM) - disable_thinking_mode = config.model.get("disable_thinking_mode", None) - thinking_mode = disable_thinking_mode is False - - # Always use ThinkingMarkerDetector for step boundary detection - if "min_step_tokens" not in config.strategy: - log.warning("strategy.min_step_tokens not set, defaulting to 50") - if "max_step_tokens" not in config.strategy: - log.warning("strategy.max_step_tokens not set, defaulting to 300") - detector = ThinkingMarkerDetector( - min_step_tokens=config.strategy.get("min_step_tokens", 50), - max_step_tokens=config.strategy.get("max_step_tokens", 300), - use_sequence=config.strategy.get("use_sequence", True), - use_conclusion=config.strategy.get("use_conclusion", True), - use_thinking=config.strategy.get("use_thinking", True), - use_verification=config.strategy.get("use_verification", True), - use_structure=config.strategy.get("use_structure", False), - use_reasoning=config.strategy.get("use_reasoning", False), - use_correction=config.strategy.get("use_correction", False), - custom_markers=config.strategy.get("custom_markers"), - ) + # All strategies use boundary detection via early stopping + from lm_polygraph.utils import APIWithUncertainty - generation_parameters = GenerationParameters() - generation_parameters.temperature = config.generation.temperature - generation_parameters.max_new_tokens = config.generation.max_new_tokens - generation_parameters.top_p = config.generation.top_p - generation_parameters.top_k = config.generation.top_k + from llm_tts.early_stopping import BoundaryEarlyStopping - # Create boundary-based early stopping - early_stopping = BoundaryEarlyStopping(detector=detector) + # Determine thinking mode (same logic as vLLM) + disable_thinking_mode = config.model.get("disable_thinking_mode", None) + thinking_mode = disable_thinking_mode is False - supports_logprobs = config.model.get("supports_logprobs", True) + # Always use ThinkingMarkerDetector for step boundary detection + if "min_step_tokens" not in config.strategy: + log.warning("strategy.min_step_tokens not set, defaulting to 50") + if "max_step_tokens" not in config.strategy: + log.warning("strategy.max_step_tokens not set, defaulting to 300") + detector = ThinkingMarkerDetector( + min_step_tokens=config.strategy.get("min_step_tokens", 50), + max_step_tokens=config.strategy.get("max_step_tokens", 300), + use_sequence=config.strategy.get("use_sequence", True), + use_conclusion=config.strategy.get("use_conclusion", True), + use_thinking=config.strategy.get("use_thinking", True), + use_verification=config.strategy.get("use_verification", True), + use_structure=config.strategy.get("use_structure", False), + use_reasoning=config.strategy.get("use_reasoning", False), + use_correction=config.strategy.get("use_correction", False), + custom_markers=config.strategy.get("custom_markers"), + ) - model = BlackboxModelWithStreaming( - openai_api_key=api_key, - model_path=model_path, - supports_logprobs=supports_logprobs, - early_stopping=early_stopping, - generation_parameters=generation_parameters, - base_url=base_url, - ) + generation_parameters = GenerationParameters() + generation_parameters.temperature = config.generation.temperature + generation_parameters.max_new_tokens = config.generation.max_new_tokens + generation_parameters.top_p = config.generation.top_p + generation_parameters.top_k = config.generation.top_k - # Set up uncertainty scorer if logprobs are supported and scorer is configured - scorer_type = config.scorer.type if config.scorer else None - if supports_logprobs and scorer_type and POLYGRAPH_UNCERTAINTY_AVAILABLE: - if scorer_type == "perplexity": - stat_calculators = [VLLMLogprobsCalculator()] - estimator = Perplexity() - elif scorer_type == "sequence_prob": - stat_calculators = [VLLMLogprobsCalculator()] - estimator = MaximumSequenceProbability() - elif scorer_type == "uncertainty_pd": - from llm_tts.scorers.estimator_uncertainty_pd import PDGap - - stat_calculators = [VLLMLogprobsCalculator(output_matrix=True)] - estimator = PDGap() - elif scorer_type == "entropy": - stat_calculators = [VLLMLogprobsCalculator(), EntropyCalculator()] - estimator = MeanTokenEntropy() - elif scorer_type == "prm": - # PRM uses its own model; use entropy for generation scoring - stat_calculators = [VLLMLogprobsCalculator(), EntropyCalculator()] - estimator = MeanTokenEntropy() - else: - stat_calculators = None - estimator = None - - if stat_calculators and estimator: - # Wrap model with uncertainty scorer (same pattern as VLLMWithUncertainty) - model = APIWithUncertainty( - model=model, - stat_calculators=stat_calculators, - estimator=estimator, - ) - log.info( - f"Wrapped model with APIWithUncertainty({type(estimator).__name__})" - ) + # Create boundary-based early stopping + early_stopping = BoundaryEarlyStopping(detector=detector) - step_generator = StepCandidateGeneratorThroughAPI( - model=model, - thinking_mode=thinking_mode, - detector=detector, - answer_patterns=config.strategy.get( - "detector_answer_patterns", - [], - ), - max_new_tokens=config.generation.max_new_tokens, - temperature=config.generation.temperature, - top_p=config.generation.top_p, - top_k=config.generation.get("top_k", 20), - presence_penalty=config.generation.get("presence_penalty", 0.0), - max_context_budget=config.model.get( - "max_context_budget", config.model.get("max_model_len", 32768) - ), - prefill_mode=config.model.get("prefill_mode", False), - disable_thinking_mode=disable_thinking_mode, - supports_logprobs=supports_logprobs, - max_concurrent_requests=config.model.get( - "max_concurrent_requests", 256 - ), - ) + supports_logprobs = config.model.get("supports_logprobs", True) - log.info( - f"Created API step generator: thinking_mode={thinking_mode}, " - f"uncertainty={'APIWithUncertainty' if hasattr(model, 'estimator') else 'no'}" - ) + model = BlackboxModelWithStreaming( + openai_api_key=api_key, + model_path=model_path, + supports_logprobs=supports_logprobs, + early_stopping=early_stopping, + generation_parameters=generation_parameters, + base_url=base_url, + ) + + # Set up uncertainty scorer if logprobs are supported and scorer is configured + scorer_type = config.scorer.type if config.scorer else None + if supports_logprobs and scorer_type and POLYGRAPH_UNCERTAINTY_AVAILABLE: + if scorer_type == "perplexity": + stat_calculators = [VLLMLogprobsCalculator()] + estimator = Perplexity() + elif scorer_type == "sequence_prob": + stat_calculators = [VLLMLogprobsCalculator()] + estimator = MaximumSequenceProbability() + elif scorer_type == "uncertainty_pd": + from llm_tts.scorers.estimator_uncertainty_pd import PDGap + + stat_calculators = [VLLMLogprobsCalculator(output_matrix=True)] + estimator = PDGap() + elif scorer_type == "entropy": + stat_calculators = [VLLMLogprobsCalculator(), EntropyCalculator()] + estimator = MeanTokenEntropy() + elif scorer_type == "prm": + # PRM uses its own model; use entropy for generation scoring + stat_calculators = [VLLMLogprobsCalculator(), EntropyCalculator()] + estimator = MeanTokenEntropy() + else: + stat_calculators = None + estimator = None + + if stat_calculators and estimator: + # Wrap model with uncertainty scorer (same pattern as VLLMWithUncertainty) + model = APIWithUncertainty( + model=model, + stat_calculators=stat_calculators, + estimator=estimator, + ) + log.info( + f"Wrapped model with APIWithUncertainty({type(estimator).__name__})" + ) + + step_generator = StepCandidateGeneratorThroughAPI( + model=model, + thinking_mode=thinking_mode, + detector=detector, + answer_patterns=config.strategy.get( + "detector_answer_patterns", + [], + ), + max_new_tokens=config.generation.max_new_tokens, + temperature=config.generation.temperature, + top_p=config.generation.top_p, + top_k=config.generation.get("top_k", 20), + presence_penalty=config.generation.get("presence_penalty", 0.0), + max_context_budget=config.model.get( + "max_context_budget", config.model.get("max_model_len", 32768) + ), + prefill_mode=config.model.get("prefill_mode", False), + disable_thinking_mode=disable_thinking_mode, + supports_logprobs=supports_logprobs, + max_concurrent_requests=config.model.get( + "max_concurrent_requests", 256 + ), + ) + + log.info( + f"Created API step generator: thinking_mode={thinking_mode}, " + f"uncertainty={'APIWithUncertainty' if hasattr(model, 'estimator') else 'no'}" + ) else: raise ValueError(f"Model type {config.model.type} not supported") @@ -888,24 +874,18 @@ def create_tts_strategy( batch_size=config.strategy.get("batch_size", 1000), ) elif config.strategy.type == "deepconf": - # DeepConf supports both API models (with logprobs) and local HuggingFace models - # Validation is done inside StrategyDeepConf.__init__ + # DeepConf uses step_generator for generation, extracts logprobs for confidence + data_name = config.strategy.get("data_name", None) strategy = StrategyDeepConf( - model=model, - mode=config.strategy.mode, + step_generator=step_generator, budget=config.strategy.get("budget", 8), - warmup_traces=config.strategy.get("warmup_traces", 4), - total_budget=config.strategy.get("total_budget", 10), - confidence_percentile=config.strategy.get("confidence_percentile", 90), window_size=config.strategy.get("window_size", 2048), filter_method=config.strategy.get("filter_method", "top10"), - temperature=config.strategy.get("temperature", 0.7), - top_p=config.strategy.get("top_p", 1.0), - max_tokens=config.strategy.get("max_tokens", 512), - top_logprobs=config.strategy.get("top_logprobs", 20), - n_threads=config.strategy.get("n_threads", 8), - disable_thinking_mode=config.model.get("disable_thinking_mode", True), - seed=config.system.seed, + confidence_threshold=config.strategy.get("confidence_threshold", None), + data_name=data_name, + mode=config.strategy.get("mode", "offline"), + online_batch_size=config.strategy.get("online_batch_size", 4), + min_agreement=config.strategy.get("min_agreement", 0.8), ) elif config.strategy.type == "beam_search":