diff --git a/.gitignore b/.gitignore index 222caa50..db4719f3 100644 --- a/.gitignore +++ b/.gitignore @@ -49,7 +49,7 @@ htmlcov/ .tox/ .coverage .coverage.* -.cache +.cache* nosetests.xml coverage.xml *.cover diff --git a/cookbook/exp/embedding/ablation_all.sh b/cookbook/exp/embedding/ablation_all.sh new file mode 100755 index 00000000..2be4c7fc --- /dev/null +++ b/cookbook/exp/embedding/ablation_all.sh @@ -0,0 +1,84 @@ +#!/bin/bash +# RAG Ablation Suite — 串行运行全部 5 个消融实验 +# GPUs: 需要 8 卡（兼容所有配置的最大需求） +# +# 用法: +# COMPRESS_API_KEY=sk-xxx bash cookbook/exp/embedding/ablation_all.sh + +set -euo pipefail + +export COMPRESS_API_KEY="${COMPRESS_API_KEY:?Set COMPRESS_API_KEY}" + +SCRIPT="cookbook/exp/embedding/eval_gpqa_rag.py" +N=500 +SEED=100 +SIM=0.6 +TOPK=1 +OUTDIR="./output/thinking_rag" +DB_PATH="./output.oldemb/thinking_rag/lance.db" + +# echo "============================================================" +# echo " Ablation 1/5: Direct (no RAG)" +# echo "============================================================" +# GEN_GPUS=8 python $SCRIPT \ +# --mode direct --n $N --seed $SEED \ +# --output $OUTDIR/ablation_direct_65k.jsonl + +# echo "" +# echo "============================================================" +# echo " Ablation 2/5: RAG + raw thinking (drop >24k, no condenser)" +# echo "============================================================" +# python $SCRIPT \ +# --mode rag --n $N --seed $SEED \ +# --db-path $DB_PATH \ +# --sim-threshold $SIM --top-k $TOPK \ +# --max-trace-len 24000 \ +# --output $OUTDIR/ablation_rag_raw_24k.jsonl + +echo "" +echo "============================================================" +echo " Ablation 3/5: RAG + API condenser (qwen3.7-max)" +echo "============================================================" +python $SCRIPT \ + --mode rag --n $N --seed $SEED \ + --db-path $DB_PATH \ + --sim-threshold $SIM --top-k $TOPK \ + --condense \ + --output $OUTDIR/ablation_rag_api_condenser_65k.jsonl + +echo "" +echo "============================================================" +echo " Ablation 4/5: RAG + local vLLM condenser (4B) + API fallback" +echo "============================================================" +EVAL_CONDENSER_GPUS=2 python $SCRIPT \ + --mode rag --n $N --seed $SEED \ + --db-path $DB_PATH \ + --sim-threshold $SIM --top-k $TOPK \ + --condense \ + --output $OUTDIR/ablation_rag_local_condenser_65k.jsonl + +# echo "" +# echo "============================================================" +# echo " Ablation 5/5: RAG + cot_compressed (pre-compressed, no runtime condenser)" +# echo "============================================================" +# python $SCRIPT \ +# --mode rag --n $N --seed $SEED \ +# --db-path $DB_PATH \ +# --sim-threshold $SIM --top-k $TOPK \ +# --use-cot-compressed \ +# --max-trace-len 4000 \ +# --output $OUTDIR/ablation_rag_cot_compressed_65k.jsonl + +echo "" +echo "============================================================" +echo " All 5 ablations complete. Results:" +echo "============================================================" +for f in $OUTDIR/ablation_*_65k.jsonl $OUTDIR/ablation_*_24k.jsonl; do + n=$(wc -l < "$f") + correct=$(python -c " +import json +recs=[json.loads(l) for l in open('$f') if l.strip()] +print(sum(1 for r in recs if r['is_correct'])) +") + echo " $(basename $f): $correct/$n" +done diff --git a/cookbook/exp/embedding/ablation_direct.sh b/cookbook/exp/embedding/ablation_direct.sh new file mode 100755 index 00000000..f5e5d145 --- /dev/null +++ b/cookbook/exp/embedding/ablation_direct.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# Ablation 1: Direct (no RAG, no condenser) +# GPUs: 4 (gen only) +# Baseline — model solves problems without any retrieved context. + +set -euo pipefail + +python cookbook/exp/embedding/eval_gpqa_rag.py \ + --mode direct \ + --n 200 \ + --seed 42 \ + --output ./output/thinking_rag/ablation_direct.jsonl diff --git a/cookbook/exp/embedding/ablation_rag_api_condenser.sh b/cookbook/exp/embedding/ablation_rag_api_condenser.sh new file mode 100755 index 00000000..fb47718e --- /dev/null +++ b/cookbook/exp/embedding/ablation_rag_api_condenser.sh @@ -0,0 +1,17 @@ +#!/bin/bash +# Ablation 3: RAG + API condenser (qwen3.7-max) +# GPUs: 6 (emb=2 + gen=4), condenser via API (no local vLLM) +# Compresses thinking_raw with COMPRESS_SYSTEM + CONDENSE_EVAL_QUERY via API. + +set -euo pipefail + +export COMPRESS_API_KEY="${COMPRESS_API_KEY:?Set COMPRESS_API_KEY}" + +python cookbook/exp/embedding/eval_gpqa_rag.py \ + --mode rag \ + --n 200 \ + --seed 42 \ + --sim-threshold 0.6 \ + --top-k 1 \ + --condense \ + --output ./output/thinking_rag/ablation_rag_api_condenser.jsonl diff --git a/cookbook/exp/embedding/ablation_rag_local_condenser.sh b/cookbook/exp/embedding/ablation_rag_local_condenser.sh new file mode 100755 index 00000000..defa863c --- /dev/null +++ b/cookbook/exp/embedding/ablation_rag_local_condenser.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Ablation 4: RAG + local vLLM condenser (Qwen3.5-4B-CM-v2) +# GPUs: 8 (emb=2 + gen=4 + condenser=2) +# Local 4B condenser as primary, API as fallback. + +set -euo pipefail + +export EVAL_CONDENSER_GPUS=2 +export COMPRESS_API_KEY="${COMPRESS_API_KEY:?Set COMPRESS_API_KEY}" + +python cookbook/exp/embedding/eval_gpqa_rag.py \ + --mode rag \ + --n 200 \ + --seed 42 \ + --sim-threshold 0.6 \ + --top-k 1 \ + --condense \ + --output ./output/thinking_rag/ablation_rag_local_condenser.jsonl diff --git a/cookbook/exp/embedding/ablation_rag_raw.sh b/cookbook/exp/embedding/ablation_rag_raw.sh new file mode 100755 index 00000000..86e35bd1 --- /dev/null +++ b/cookbook/exp/embedding/ablation_rag_raw.sh @@ -0,0 +1,15 @@ +#!/bin/bash +# Ablation 2: RAG + raw thinking (no condenser, truncated to max-trace-len) +# GPUs: 6 (emb=2 + gen=4) +# Uses thinking_raw directly, truncated to 4000 chars. + +set -euo pipefail + +python cookbook/exp/embedding/eval_gpqa_rag.py \ + --mode rag \ + --n 200 \ + --seed 42 \ + --sim-threshold 0.6 \ + --top-k 1 \ + --max-trace-len 4000 \ + --output ./output/thinking_rag/ablation_rag_raw.jsonl diff --git a/cookbook/exp/embedding/build_thinking_rag_index.py b/cookbook/exp/embedding/build_thinking_rag_index.py index d228a597..a71bae06 100644 --- a/cookbook/exp/embedding/build_thinking_rag_index.py +++ b/cookbook/exp/embedding/build_thinking_rag_index.py @@ -39,6 +39,8 @@ import os import re import sys +import threading +import time from pathlib import Path from typing import Any, Dict, Iterator, List, Optional, Tuple @@ -163,7 +165,7 @@ EMBED_MODEL_ID = os.environ.get( 'EMBED_MODEL_ID', - 'output/embedding_lora_transformers/step_4000', + 'output/embedding_full_transformers/last-checkpoint', ) CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') @@ -185,30 +187,42 @@ SIM_THRESHOLD = float(os.environ.get('SIM_THRESHOLD', 0.65)) MIN_TEXT_CHARS = int(os.environ.get('MIN_TEXT_CHARS', 256)) +# Dataset mix caps (only used in 'both' mode). None = no cap. +THINK_CAP: Optional[int] = int(os.environ.get('THINK_CAP', 400_000)) or None +INDEX_CAP: Optional[int] = int(os.environ.get('INDEX_CAP', 400_000)) or None +MIX_SHUFFLE_SEED = 100 + +# Concurrency knobs for API fallback and prefetch pipeline. +API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 8)) +API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) +PREFETCH_WORKERS = int(os.environ.get('PREFETCH_WORKERS', 2)) + # Hard-templated hints: the condenser SFT prior maps `Skill` to the legacy # `Use when: / numbered steps / Output:` skeleton on long inputs; embedding the # exact 4-line body template + explicit negative constraints is the only way to # override it deterministically across query and cot sides. RAG_QUERY_HINT = ( - 'Summarize this query for retrieval. ' - 'The body of ## Summary MUST follow this EXACT 4-line template — ' - 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' + 'Extract the abstract PROBLEM TYPE from this query. ' + 'IGNORE all specific numbers, values, variable names, and parameters — ' + 'focus ONLY on the CLASS of problem and the METHODOLOGY required. ' + 'The body of ## Summary MUST follow this EXACT 4-line template:\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the specific pattern, never generic labels.') + 'Topic must name the method class, never mention specific numbers.') RAG_THINKING_HINT = ( - 'Summarize this reasoning trace for retrieval. ' - 'The body of ## Summary MUST follow this EXACT 4-line template — ' - 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' + 'Extract the abstract METHODOLOGY demonstrated in this solution. ' + 'IGNORE all specific numbers, values, and computed results — ' + 'focus ONLY on the general TECHNIQUE and key reasoning STEPS. ' + 'The body of ## Summary MUST follow this EXACT 4-line template:\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the specific pattern, never generic labels.') + 'Topic must name the method class, never mention specific numbers.') # OpenAI API fallback (used when vLLM truncates). COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '') @@ -419,23 +433,124 @@ def _api_compress(api: OpenAIClient, messages: List[Dict[str, str]]) -> Optional return _strip_outer_codefence(content) +_api_throttle_lock = threading.Lock() +_api_last_call = [0.0] + + +def _api_throttle(): + with _api_throttle_lock: + gap = time.monotonic() - _api_last_call[0] + if gap < API_MIN_INTERVAL: + time.sleep(API_MIN_INTERVAL - gap) + _api_last_call[0] = time.monotonic() + + +def _api_compress_throttled(api: OpenAIClient, messages: List[Dict[str, str]]) -> Optional[str]: + """Rate-limited API compression call.""" + _api_throttle() + return _api_compress(api, messages) + + def _resolve_compressed(sampler: vLLMSampler, api: Optional[OpenAIClient], texts: List[str], query_hint: str) -> List[Optional[str]]: - """Run vLLM batch; replace truncations / skeleton-incomplete with API output.""" + """Run vLLM batch; replace truncations / skeleton-incomplete with API output. + + API fallback runs concurrently (up to API_CONCURRENCY workers) for speed. + """ pairs = _vllm_compress(sampler, texts, query_hint) - results: List[Optional[str]] = [] - for (text, stop), src_text in zip(pairs, texts): + results: List[Optional[str]] = [None] * len(texts) + fallback_indices: List[int] = [] + for i, ((text, stop), src_text) in enumerate(zip(pairs, texts)): if stop != 'length' and not _is_truncated_compression(text): - results.append(text) - continue - if api is None: - results.append(None) + results[i] = text + else: + fallback_indices.append(i) + + if fallback_indices and api is not None: + from concurrent.futures import ThreadPoolExecutor, as_completed + with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: + futures = {} + for idx in fallback_indices: + msgs = _build_compress_messages(texts[idx], query_hint) + futures[pool.submit(_api_compress_throttled, api, msgs)] = idx + for fut in as_completed(futures): + idx = futures[fut] + api_text = fut.result() + if api_text and not _is_truncated_compression(api_text): + results[idx] = api_text + + return results + + +def _resolve_compressed_multi(sampler: vLLMSampler, api: Optional[OpenAIClient], + texts: List[str], hints: List[str]) -> List[Optional[str]]: + """Like _resolve_compressed but each text has its own per-item hint. + + Merges all texts into a SINGLE vLLM batch call (instead of one per hint), + dramatically reducing round-trip overhead when processing interleaved + query+cot pairs with different hint strings. + + Args: + sampler: vLLM condenser sampler. + api: Optional OpenAI-compatible API client for fallback. + texts: List of raw texts to compress (may contain empty strings to skip). + hints: Per-text hint strings (same length as texts). + + Returns: + List of compressed texts (None where compression failed entirely). + """ + assert len(texts) == len(hints), f'texts({len(texts)}) != hints({len(hints)})' + if not texts: + return [] + + # Skip texts that would exceed the condenser's context window. + _max_input_chars = (CONDENSE_MAX_MODEL_LEN - CONDENSE_MAX_TOKENS) * 3 + skip_mask = [len(t) > _max_input_chars for t in texts] + + # Build prompts per-item (each text gets its own hint as the query parameter). + prompts = [{'messages': _build_compress_messages(t, h)} + for t, h, skip in zip(texts, hints, skip_mask) if not skip] + active_indices = [i for i, skip in enumerate(skip_mask) if not skip] + params = TwinkleSamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, + temperature=COMPRESS_TEMPERATURE, + top_p=COMPRESS_TOP_P, + num_samples=1, + ) + + # Single vLLM batch call — the key throughput win. + responses = sampler.sample(prompts, params) if prompts else [] + + results: List[Optional[str]] = [None] * len(texts) + fallback_indices: List[int] = [] + for resp_idx, orig_idx in enumerate(active_indices): + resp = responses[resp_idx] + seq = resp.sequences[0] if resp and resp.sequences else None + if seq is None: + fallback_indices.append(orig_idx) continue - api_text = _api_compress(api, _build_compress_messages(src_text, query_hint)) - if api_text is None or _is_truncated_compression(api_text): - results.append(None) + text = seq.decoded or '' + text = re.sub(r'<\|[^|]+\|>', '', text).rstrip() + text = _strip_outer_codefence(text) + if seq.stop_reason != 'length' and not _is_truncated_compression(text): + results[orig_idx] = text else: - results.append(api_text) + fallback_indices.append(orig_idx) + + # Concurrent API fallback for failed items. + if fallback_indices and api is not None: + from concurrent.futures import ThreadPoolExecutor, as_completed + with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: + futures = {} + for idx in fallback_indices: + msgs = _build_compress_messages(texts[idx], hints[idx]) + futures[pool.submit(_api_compress_throttled, api, msgs)] = idx + for fut in as_completed(futures): + idx = futures[fut] + api_text = fut.result() + if api_text and not _is_truncated_compression(api_text): + results[idx] = api_text + return results @@ -545,7 +660,7 @@ def _existing_ids(table) -> set: def _stream_corpus(total: Optional[int], load_from_cache_file: bool, max_rows: int = 0) -> Iterator[Dict[str, Any]]: - ds = _GET_DATASET(total=total, load_from_cache_file=load_from_cache_file) + ds = _GET_DATASET(total=total or None, load_from_cache_file=load_from_cache_file) n_full = len(ds) cap = max_rows if (max_rows and max_rows < n_full) else n_full sys.stderr.write(f'[corpus] get_dataset: {n_full} rows' @@ -602,32 +717,58 @@ def build_index(args: argparse.Namespace, # ---- Streaming loop ----------------------------------------------------- n_seen = n_kept = n_dropped_short = n_dropped_compress = n_dropped_sim = 0 n_dropped_dup = 0 - pbar = tqdm(desc='index', unit='row', dynamic_ncols=True) + n_no_id = 0 + n_no_query = 0 + n_short_cot = 0 + _diag_samples = 5 # print first N dropped rows for diagnosis batch: List[Dict[str, Any]] = [] - def _flush(rows: List[Dict[str, Any]]) -> None: - nonlocal n_kept, n_dropped_compress, n_dropped_sim + def _compress_batch(rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Phase 1: compress query+cot in a SINGLE merged vLLM call for throughput.""" if not rows: - return - # Phase 1 — compress query (RAG_QUERY_HINT) and cot (RAG_THINKING_HINT). - # Short queries bypass condenser (passthrough) — matches training behaviour. - long_q_indices = [i for i, r in enumerate(rows) if len(r['query_raw']) >= MIN_TEXT_CHARS] - q_compressed: List[Optional[str]] = [None] * len(rows) - for i, r in enumerate(rows): - if len(r['query_raw']) < MIN_TEXT_CHARS: - q_compressed[i] = r['query_raw'] - if long_q_indices: - long_results = _resolve_compressed( - sampler, api, [rows[i]['query_raw'] for i in long_q_indices], RAG_QUERY_HINT) - for idx, res in zip(long_q_indices, long_results): - q_compressed[idx] = res - c_compressed = _resolve_compressed( - sampler, api, [r['cot_raw'] for r in rows], RAG_THINKING_HINT) + return [] + # Build a merged prompt list: interleave query and cot texts so the sampler + # processes both in one round-trip instead of two serial calls. + all_texts: List[str] = [] + all_hints: List[str] = [] + passthrough_map: Dict[int, str] = {} # prompt_idx → raw text for short queries + for r in rows: + q_raw = r['query_raw'] + if len(q_raw) < MIN_TEXT_CHARS: + passthrough_map[len(all_texts)] = q_raw + all_texts.append('') # placeholder + all_hints.append(RAG_QUERY_HINT) + else: + all_texts.append(q_raw) + all_hints.append(RAG_QUERY_HINT) + all_texts.append(r['cot_raw']) + all_hints.append(RAG_THINKING_HINT) + + # Split into passthrough vs sampler-needed + sampler_indices = [i for i in range(len(all_texts)) if i not in passthrough_map] + sampler_texts = [all_texts[i] for i in sampler_indices] + sampler_hints = [all_hints[i] for i in sampler_indices] + + # Single merged vLLM call — group by hint to maximize prefix-sharing + # (both hints produce the same COMPRESS_SYSTEM, so batching is efficient). + sampler_results = _resolve_compressed_multi( + sampler, api, sampler_texts, sampler_hints) + + # Reassemble full results + all_results: List[Optional[str]] = [None] * len(all_texts) + for idx, text in passthrough_map.items(): + all_results[idx] = text + for pos, res in zip(sampler_indices, sampler_results): + all_results[pos] = res + + # Pair up (query, cot) and filter kept_rows: List[Dict[str, Any]] = [] - for r, q_cmp, c_cmp in zip(rows, q_compressed, c_compressed): + for i, r in enumerate(rows): + q_cmp = all_results[i * 2] + c_cmp = all_results[i * 2 + 1] if not q_cmp or not c_cmp: - n_dropped_compress += 1 + nonlocal_counters['n_dropped_compress'] += 1 _log_miss(misses_path, misses_lock, { 'id': r['id'], 'source': r['source'], 'reason': 'compress_fail', 'query_raw_head': _short(r['query_raw'], 200), @@ -637,15 +778,17 @@ def _flush(rows: List[Dict[str, Any]]) -> None: r['query_compressed'] = q_cmp r['cot_compressed'] = c_cmp kept_rows.append(r) + return kept_rows + + def _embed_and_insert(kept_rows: List[Dict[str, Any]]) -> None: + """Phase 2+3: embed compressed texts and insert into LanceDB.""" if not kept_rows: return - # Phase 2 — encode anchor (compressed query) + positive (compressed cot). anchor_emb = get_embeddings( emb_model, emb_template, [r['query_compressed'] for r in kept_rows], role='anchor') positive_emb = get_embeddings( emb_model, emb_template, [r['cot_compressed'] for r in kept_rows], role='positive') sims = (anchor_emb * positive_emb).sum(axis=1).astype(np.float32) - # Phase 3 — sim filter + LanceDB insert. to_insert: List[Dict[str, Any]] = [] for idx, (r, sim_val) in enumerate(zip(kept_rows, sims)): tag = 'KEEP' if sim_val >= SIM_THRESHOLD else 'DROP' @@ -653,7 +796,7 @@ def _flush(rows: List[Dict[str, Any]]) -> None: f'q={_short(r["query_raw"], 60)!r} ' f'cot={_short(r["cot_raw"], 60)!r}', flush=True) if sim_val < SIM_THRESHOLD: - n_dropped_sim += 1 + nonlocal_counters['n_dropped_sim'] += 1 _log_miss(misses_path, misses_lock, { 'id': r['id'], 'source': r['source'], 'reason': 'sim_low', 'sim': float(sim_val), @@ -677,24 +820,60 @@ def _flush(rows: List[Dict[str, Any]]) -> None: }) if to_insert: tbl.add(to_insert) - n_kept += len(to_insert) + nonlocal_counters['n_kept'] += len(to_insert) indexed.update(r['id'] for r in to_insert) + def _process_batch(rows: List[Dict[str, Any]]) -> None: + """Full pipeline for one batch: compress → embed → insert.""" + kept = _compress_batch(rows) + _embed_and_insert(kept) + + # Mutable counters shared with nested functions (avoid nonlocal limitation). + nonlocal_counters = { + 'n_kept': 0, 'n_dropped_compress': 0, 'n_dropped_sim': 0, + } + + from concurrent.futures import ThreadPoolExecutor as _PrefetchPool + prefetch_pool = _PrefetchPool(max_workers=PREFETCH_WORKERS) + try: + # Phase 1: Stream corpus, filter rows, collect batches (fast). + pending_futures = [] + sys.stderr.write('[build] streaming corpus and submitting batches...\n') + for row in _stream_corpus(total=args.total, load_from_cache_file=not args.no_cache, max_rows=args.max_rows): n_seen += 1 - if args.limit and n_kept >= args.limit: + if args.limit and nonlocal_counters['n_kept'] >= args.limit: break rid = row.get('id') or '' if not rid: + n_no_id += 1 + if n_no_id <= _diag_samples: + sys.stderr.write(f'[diag:no_id] row keys={list(row.keys())}\n') continue if rid in indexed: n_dropped_dup += 1 continue user_query, cot = _extract_query_cot(row) - if not user_query or len(cot) < MIN_TEXT_CHARS: + if not user_query: + n_no_query += 1 + n_dropped_short += 1 + if n_no_query <= _diag_samples: + msgs = row.get('messages') + sys.stderr.write( + f'[diag:no_query] id={rid} source={row.get("source","?")} ' + f'msgs_type={type(msgs).__name__} ' + f'msgs_len={len(msgs) if isinstance(msgs, list) else "?"} ' + f'msg0_keys={list(msgs[0].keys()) if isinstance(msgs, list) and msgs and isinstance(msgs[0], dict) else "?"}\n') + continue + if len(cot) < MIN_TEXT_CHARS: + n_short_cot += 1 n_dropped_short += 1 + if n_short_cot <= _diag_samples: + sys.stderr.write( + f'[diag:short_cot] id={rid} source={row.get("source","?")} ' + f'cot_len={len(cot)} query_len={len(user_query)}\n') continue batch.append({ 'id': rid, @@ -703,21 +882,45 @@ def _flush(rows: List[Dict[str, Any]]) -> None: 'cot_raw': cot, }) if len(batch) >= args.batch_size: - _flush(batch) + pending_futures.append(prefetch_pool.submit(_process_batch, list(batch))) batch.clear() - pbar.set_postfix(kept=n_kept, sim_drop=n_dropped_sim, - cmp_drop=n_dropped_compress, refresh=False) - pbar.update(1) + + # Flush remainder if batch: - _flush(batch) + pending_futures.append(prefetch_pool.submit(_process_batch, list(batch))) batch.clear() + + n_batches = len(pending_futures) + n_valid = n_seen - n_no_id - n_dropped_dup - n_dropped_short + sys.stderr.write( + f'[build] stream done: seen={n_seen} valid={n_valid} ' + f'batches={n_batches} (no_id={n_no_id} no_query={n_no_query} ' + f'short_cot={n_short_cot} dup={n_dropped_dup})\n') + + # Phase 2: Wait for all futures with real progress tracking. + pbar = tqdm(total=n_batches, desc='compress+embed', unit='batch', + dynamic_ncols=True) + for fut in pending_futures: + fut.result() + n_kept = nonlocal_counters['n_kept'] + n_dropped_sim = nonlocal_counters['n_dropped_sim'] + n_dropped_compress = nonlocal_counters['n_dropped_compress'] + pbar.set_postfix(kept=n_kept, sim_drop=n_dropped_sim, + cmp_drop=n_dropped_compress, refresh=False) + pbar.update(1) finally: pbar.close() + prefetch_pool.shutdown(wait=True) + + n_kept = nonlocal_counters['n_kept'] + n_dropped_sim = nonlocal_counters['n_dropped_sim'] + n_dropped_compress = nonlocal_counters['n_dropped_compress'] sys.stderr.write( - f'[build] seen={n_seen} kept={n_kept} sim_drop={n_dropped_sim} ' - f'cmp_drop={n_dropped_compress} short_drop={n_dropped_short} ' - f'dup_skip={n_dropped_dup}\n') + f'[build] summary: seen={n_seen} kept={n_kept} ' + f'dup={n_dropped_dup} no_id={n_no_id} no_query={n_no_query} ' + f'short_cot={n_short_cot} compress_fail={n_dropped_compress} ' + f'sim_drop={n_dropped_sim}\n') # ---- Build vector index for fast retrieval ------------------------------ if n_kept >= 64 and not args.skip_index: @@ -864,17 +1067,17 @@ def parse_args() -> argparse.Namespace: help='LanceDB table name within --db-path.') p.add_argument('--total', type=int, default=0, help='Total dataset rows to scale corpus to (0 = base sizes from the loader module).') - p.add_argument('--dataset-module', default='dataset_index', - choices=['dataset_index', 'dataset_think'], - help='Which loader to use: dataset_index (RAG profile) or ' - 'dataset_think (training mix).') + p.add_argument('--dataset-module', default='both', + choices=['dataset_index', 'dataset_think', 'both'], + help='Which loader to use: dataset_index (RAG profile), ' + 'dataset_think (training mix), or both (50/50 mix).') p.add_argument('--limit', type=int, default=0, help='Stop building once this many rows are kept (0 = no cap).') p.add_argument('--max-rows', type=int, default=0, help='Truncate corpus to this many rows AFTER get_dataset (0 = no cap). ' 'Use this instead of --total to avoid invalidating the dataset cache.') - p.add_argument('--batch-size', type=int, default=64, - help='Rows per condense+encode batch.') + p.add_argument('--batch-size', type=int, default=128, + help='Rows per condense+encode batch (larger = better GPU util).') p.add_argument('--no-cache', action='store_true', help='Disable load_from_cache_file in dataset_think.get_dataset.') p.add_argument('--overwrite', action='store_true', @@ -901,6 +1104,27 @@ def main() -> None: if args.dataset_module == 'dataset_think': from dataset_think import get_dataset as _swap _GET_DATASET = _swap + elif args.dataset_module == 'both': + from dataset_think import get_dataset as _get_think + from datasets import concatenate_datasets + + def _get_both(total=None, load_from_cache_file=True, **kw): + _total = total or None # CLI default 0 means "no scaling" → None + ds_index = _default_get_dataset(total=_total, load_from_cache_file=load_from_cache_file) + ds_think = _get_think(total=_total, load_from_cache_file=load_from_cache_file) + if INDEX_CAP and len(ds_index.dataset) > INDEX_CAP: + ds_index.dataset = ds_index.dataset.select(range(INDEX_CAP)) + if THINK_CAP and len(ds_think.dataset) > THINK_CAP: + ds_think.dataset = ds_think.dataset.select(range(THINK_CAP)) + n_index = len(ds_index.dataset) + n_think = len(ds_think.dataset) + ds_index.dataset = concatenate_datasets( + [ds_index.dataset, ds_think.dataset]).shuffle(seed=MIX_SHUFFLE_SEED) + sys.stderr.write(f'[mix] index={n_index} + think={n_think} ' + f'→ total={len(ds_index.dataset)}\n') + return ds_index + + _GET_DATASET = _get_both sys.stderr.write(f'[main] dataset loader: {args.dataset_module}\n') # Build/eval both depend on the same Twinkle stack — initialize once. diff --git a/cookbook/exp/embedding/compare_math_levels.py b/cookbook/exp/embedding/compare_math_levels.py new file mode 100644 index 00000000..d488909b --- /dev/null +++ b/cookbook/exp/embedding/compare_math_levels.py @@ -0,0 +1,91 @@ +"""Compare MATH direct vs RAG by difficulty level. + +Re-grades both result files with the production ``answers_match`` (so the +stored ``is_correct`` is never trusted) and prints the per-level accuracy +plus the RAG gain (delta) so you can see how it varies with difficulty. + +Defaults to the raw-RAG output (``math_rag_results.jsonl``); pass a second +arg to compare a different rag file (e.g. ``math_rag_hint_results.jsonl``). + +Usage: + python cookbook/exp/embedding/compare_math_levels.py \ + [direct.jsonl] [rag.jsonl] +""" +import importlib.util +import json +import os +import sys +from collections import defaultdict + +_HERE = os.path.dirname(os.path.abspath(__file__)) + + +def _load_grader(): + spec = importlib.util.spec_from_file_location( + 'egr', os.path.join(_HERE, 'eval_gpqa_rag.py')) + egr = importlib.util.module_from_spec(spec) + spec.loader.exec_module(egr) + return egr.answers_match + + +def _load(path): + return {json.loads(l)['idx']: json.loads(l) + for l in open(path, encoding='utf-8') if l.strip()} + + +def main(): + direct_path = sys.argv[1] if len(sys.argv) > 1 else \ + './output/thinking_rag/math_direct_results.jsonl' + hint_path = sys.argv[2] if len(sys.argv) > 2 else \ + './output/thinking_rag/math_rag_results.jsonl' + + answers_match = _load_grader() + D = _load(direct_path) + H = _load(hint_path) + common = sorted(set(D) & set(H)) + print(f'direct={len(D)} rag+hint={len(H)} common={len(common)}') + + def runaway(rec): + mo = rec.get('model_output') or '' + return ('' not in mo) or ( + not (rec.get('predicted') or '').strip() and len(mo) > 40000) + + def correct(rec): + return answers_match(rec.get('predicted') or '', + rec['reference_answer']) + + # level -> counters + per = defaultdict(lambda: {'n': 0, 'd': 0, 'h': 0, + 'd_run': 0, 'h_run': 0}) + for i in common: + lv = H[i].get('level') or D[i].get('level') or 'Unknown' + c = per[lv] + c['n'] += 1 + c['d'] += int(correct(D[i])) + c['h'] += int(correct(H[i])) + c['d_run'] += int(runaway(D[i])) + c['h_run'] += int(runaway(H[i])) + + print(f'\n{"level":>10} | {"n":>4} | {"direct":>7} | {"rag+hint":>8} | ' + f'{"delta":>7} | {"d_run":>6} | {"h_run":>6}') + print('-' * 68) + tot = {'n': 0, 'd': 0, 'h': 0, 'd_run': 0, 'h_run': 0} + for lv in sorted(per.keys()): + c = per[lv] + for k in tot: + tot[k] += c[k] + n = c['n'] + dacc, hacc = c['d'] / n, c['h'] / n + print(f'{lv:>10} | {n:>4} | {dacc:>7.3f} | {hacc:>8.3f} | ' + f'{hacc - dacc:>+7.3f} | {c["d_run"]/n:>6.1%} | ' + f'{c["h_run"]/n:>6.1%}') + print('-' * 68) + n = tot['n'] + if n: + print(f'{"OVERALL":>10} | {n:>4} | {tot["d"]/n:>7.3f} | ' + f'{tot["h"]/n:>8.3f} | {(tot["h"]-tot["d"])/n:>+7.3f} | ' + f'{tot["d_run"]/n:>6.1%} | {tot["h_run"]/n:>6.1%}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/dataset_hard.py b/cookbook/exp/embedding/dataset_hard.py new file mode 100644 index 00000000..9fa059b9 --- /dev/null +++ b/cookbook/exp/embedding/dataset_hard.py @@ -0,0 +1,202 @@ +"""Hard-negative dataset for embedding training. + +Provides ReasonIR (AI-ModelScope/reasonir-data, hq subset): + - query: reasoning-intensive question + - positive: BRIGHT document (resolved via xlangai/BRIGHT documents corpus) + - negatives: plausibly related but ultimately unhelpful documents + +Output schema: ``{id, source, query, cot, response, negatives}`` +where ``negatives`` is a list of strings (each a separate hard negative). +""" +import hashlib +import os +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional + +from datasets import Dataset as HFDataset +from modelscope import MsDataset + +_CACHE_DIR = Path(__file__).resolve().parent / '.cache_hard' + + +def _hash_id(prefix: str, content: str) -> str: + return f'{prefix}__{hashlib.md5(content.encode("utf-8")).hexdigest()[:16]}' + + +# --------------------------------------------------------------------------- +# BRIGHT document corpus (lazy singleton) +# --------------------------------------------------------------------------- + +_BRIGHT_SPLITS = [ + 'aops', 'biology', 'earth_science', 'economics', 'leetcode', 'pony', + 'psychology', 'robotics', 'stackoverflow', 'sustainable_living', + 'theoremqa_questions', 'theoremqa_theorems', +] + +_bright_docs: Optional[Dict[str, str]] = None + + +def _load_bright_docs() -> Dict[str, str]: + """Load all BRIGHT document splits into {id -> content} lookup dict.""" + global _bright_docs + if _bright_docs is not None: + return _bright_docs + sys.stderr.write('[dataset_hard] Loading BRIGHT documents corpus...\n') + _bright_docs = {} + for split in _BRIGHT_SPLITS: + try: + ds = MsDataset.load( + 'xlangai/BRIGHT', subset_name='documents', split=split, + download_mode='reuse_dataset_if_exists') + for row in ds: + doc_id = row.get('id', '') + content = row.get('content', '') + if doc_id and content: + _bright_docs[doc_id] = content + short = doc_id.rsplit('/', 1)[-1] if '/' in doc_id else doc_id + if short not in _bright_docs: + _bright_docs[short] = content + sys.stderr.write(f' [{split}] loaded {len(ds)} docs\n') + except Exception as e: + sys.stderr.write(f' [{split}] FAILED: {e}\n') + sys.stderr.write(f'[dataset_hard] BRIGHT total: {len(_bright_docs)} entries\n') + return _bright_docs + + +# --------------------------------------------------------------------------- +# ReasonIR dataset +# --------------------------------------------------------------------------- + +def get_dataset_reasonir(max_rows: Optional[int] = None, + max_negatives: int = 16, + load_from_cache_file: bool = True) -> HFDataset: + """Load AI-ModelScope/reasonir-data (hq subset) with BRIGHT doc resolution. + + Schema: {id, source, query, cot, response, negatives} + """ + cache_key = f'reasonir_neg{max_negatives}' + cache_path = _CACHE_DIR / cache_key + if load_from_cache_file and cache_path.exists(): + sys.stderr.write(f'[reasonir] loading from cache: {cache_path}\n') + ds = HFDataset.load_from_disk(str(cache_path)) + if max_rows and len(ds) > max_rows: + ds = ds.select(range(max_rows)) + sys.stderr.write(f'[reasonir] {len(ds)} rows (cached)\n') + return ds + + ds = MsDataset.load( + 'AI-ModelScope/reasonir-data', subset_name='hq', split='train', + download_mode='reuse_dataset_if_exists') + if max_rows and len(ds) > max_rows: + ds = ds.select(range(max_rows)) + + bright = _load_bright_docs() + rows = [] + n_miss = 0 + for row in ds: + query_parts = row.get('query', []) + if not isinstance(query_parts, list) or len(query_parts) < 2: + continue + query = query_parts[1].strip() + if not query: + continue + + pos_list = row.get('pos', []) + if not pos_list: + continue + pos_id = pos_list[0][1] if isinstance(pos_list[0], list) and len(pos_list[0]) > 1 else '' + cot = bright.get(pos_id, '') + if not cot: + n_miss += 1 + continue + + neg_list = row.get('neg', []) + negatives = [] + for neg in neg_list: + if isinstance(neg, list) and len(neg) > 1: + neg_text = neg[1].strip() + if neg_text: + negatives.append(neg_text) + if len(negatives) >= max_negatives: + break + + if not negatives: + continue + + rows.append({ + 'id': _hash_id('reasonir', f'{query}\n{pos_id}'), + 'source': 'reasonir-hq', + 'query': query, + 'cot': cot, + 'response': '', + 'negatives': negatives, + }) + + if n_miss: + sys.stderr.write(f'[reasonir] {n_miss} rows skipped (BRIGHT doc not found)\n') + sys.stderr.write(f'[reasonir] {len(rows)} rows with hard negatives\n') + result = HFDataset.from_dict(_rows_to_cols(rows)) + # Persist full dataset; max_rows is applied post-cache for flexibility. + cache_path.parent.mkdir(parents=True, exist_ok=True) + result.save_to_disk(str(cache_path)) + sys.stderr.write(f'[reasonir] cached to {cache_path}\n') + if max_rows and len(result) > max_rows: + result = result.select(range(max_rows)) + return result + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _rows_to_cols(rows: List[Dict[str, Any]]) -> Dict[str, list]: + if not rows: + return {'id': [], 'source': [], 'query': [], 'cot': [], + 'response': [], 'negatives': []} + keys = rows[0].keys() + return {k: [r[k] for r in rows] for k in keys} + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def get_dataset( + reasonir_max: Optional[int] = None, + max_negatives: int = 16, + load_from_cache_file: bool = True, + **kwargs, +) -> HFDataset: + """Load hard-negative dataset (reasonir only). + + Returns HF Dataset with schema: {id, source, query, cot, response, negatives} + """ + ds = get_dataset_reasonir(max_rows=reasonir_max, max_negatives=max_negatives, + load_from_cache_file=load_from_cache_file) + if len(ds) == 0: + sys.stderr.write('[dataset_hard] WARNING: reasonir dataset empty\n') + else: + sys.stderr.write(f'[dataset_hard] reasonir={len(ds)}\n') + return ds + + +if __name__ == '__main__': + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('--reasonir-max', type=int, default=1000) + args = parser.parse_args() + + ds = get_dataset(reasonir_max=args.reasonir_max) + print(f'Total rows: {len(ds)}') + print(f'Features: {ds.features}') + if len(ds) > 0: + row = ds[0] + print(f'\nSample[0]:') + print(f' id: {row["id"]}') + print(f' source: {row["source"]}') + print(f' query: {row["query"][:100]}...') + print(f' cot: {row["cot"][:100]}...') + print(f' negatives: {len(row["negatives"])} items') + if row['negatives']: + print(f' [0]: {row["negatives"][0][:80]}...') diff --git a/cookbook/exp/embedding/dataset_index.py b/cookbook/exp/embedding/dataset_index.py index c86e1c52..7d2905a5 100644 --- a/cookbook/exp/embedding/dataset_index.py +++ b/cookbook/exp/embedding/dataset_index.py @@ -704,7 +704,7 @@ def get_dataset(total: Optional[int] = None, ], dropped_log_path=dropped_log or '', ) - dataset.map(qp, batched=True, num_proc=32, load_from_cache_file=load_from_cache_file) + dataset.map(qp, num_proc=32, load_from_cache_file=load_from_cache_file) return dataset diff --git a/cookbook/exp/embedding/eval_gpqa_rag.py b/cookbook/exp/embedding/eval_gpqa_rag.py new file mode 100644 index 00000000..dd71590c --- /dev/null +++ b/cookbook/exp/embedding/eval_gpqa_rag.py @@ -0,0 +1,1490 @@ +"""Math evaluation: direct vs RAG-augmented with Qwen3.5-4B. + +Datasets (``--dataset``): + - ``math`` (default): MATH (Hendrycks), stratified by difficulty (Level 1-5) + so RAG gain can be plotted against difficulty. + - ``aops``: AoPS competition problems (metadata.boxed only). + +Modes (``--mode``): + - ``direct``: The model solves problems directly (4 GPUs, TP=4). + - ``rag`` (default): Retrieve top-k thinking traces from LanceDB, condense + them (API qwen3.7-max), inject as 1-shot examples, then solve + (8 GPUs: DP=4 embedding + TP=4 vLLM). + +Defaults implement **raw RAG on MATH**: ``--dataset math --mode rag --condense`` +with hint filtering OFF. The API condenser needs ``COMPRESS_API_KEY`` (or a +local condenser via ``EVAL_CONDENSER_GPUS``); otherwise pass ``--no-condense``. + +Optional ``--hint`` flag (rag mode only): + After retrieval + condensing, call an API model to filter and refine the + traces — keeping only applicable methods — then inject the refined trace. + +Reference answers are the ``\\boxed{...}`` content of each solution. + +Launch examples: + # Default: raw RAG on MATH, stratified 100/level (needs COMPRESS_API_KEY) + COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py + + # Paired direct baseline on the same MATH subset + python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct + + # Raw RAG without condenser (inject raw retrieved traces) + python cookbook/exp/embedding/eval_gpqa_rag.py --no-condense + + # Add hint filtering back on top of condensing + COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py --hint + + # Fall back to the old AoPS dataset + python cookbook/exp/embedding/eval_gpqa_rag.py --dataset aops +""" +import argparse +import json +import os +import random +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Dict, List, Optional + +import numpy as np +import torch + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams as TwinkleSamplingParams +from twinkle.loss import InfonceLoss +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient + +logger = get_logger() + +# -- Condenser config ---------------------------------------------------------- +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +CONDENSE_API_KEY = os.environ.get('COMPRESS_API_KEY', '') +CONDENSE_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') +CONDENSE_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') +CONDENSE_API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 32)) +CONDENSE_API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) +CONDENSE_TEMPERATURE = 0.2 +CONDENSE_MAX_TOKENS = 8192 + +# -- Hint analysis config ------------------------------------------------------ +HINT_ANALYSIS_MAX_TOKENS = int(os.environ.get('HINT_ANALYSIS_MAX_TOKENS', 2000)) +HINT_ANALYSIS_TEMPERATURE = 0.2 + +HINT_ANALYSIS_SYSTEM = ( + 'You are a mathematical reasoning trace filter. ' + 'Given a target problem and reasoning traces retrieved from SIMILAR (but different) problems, ' + 'your task is to FILTER and REFINE the traces into a clean reference.\n\n' + 'Rules:\n' + '1. KEEP: solution steps, methods, formulas, techniques, and key insights ' + 'that are directly applicable to solving the target problem.\n' + '2. REMOVE: problem-specific numeric calculations that do not transfer, ' + 'dead-end explorations, irrelevant approaches, verbose restatements, ' + 'and any content that would mislead the solver on the target problem.\n' + '3. Output the refined trace directly as actionable solution steps. ' + 'Preserve the original mathematical expressions and step structure.\n' + '4. Do NOT solve the target problem. Do NOT add your own solutions or commentary.\n' + '5. Do NOT output the answer to either problem.\n' + '6. If the traces are entirely irrelevant, output exactly: "No applicable methods."' +) + +HINT_ANALYSIS_USER = ( + '## Target Problem\n{query}\n\n' + '## Retrieved Reasoning Traces\n{thinking}' +) + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- +# -- Gen/Embed config --------------------------------------------------------- +GEN_MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3.5-4B') +EMBED_MODEL_ID = os.environ.get( + 'EMBED_MODEL_ID', 'output.oldemb/embedding_full_transformers/last-checkpoint') + +GEN_GPUS = int(os.environ.get('GEN_GPUS', 4)) +EMB_GPUS = int(os.environ.get('EMB_GPUS', 2)) +EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 20000)) + +GEN_GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.85)) +GEN_MAX_MODEL_LEN = int(os.environ.get('GEN_MAX_MODEL_LEN', 65536)) +GEN_MAX_TOKENS = int(os.environ.get('GEN_MAX_TOKENS', 65536)) +GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) +GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) + +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +MATH_DATA_DIR = os.environ.get('MATH_DATA_DIR', './output/math_data/MATH') + + +# --------------------------------------------------------------------------- +# Condenser prompts & validation +# --------------------------------------------------------------------------- + +COMPRESS_SYSTEM = """\ +You are a reasoning-trace condenser. Given a verbose reasoning trace, \ +extract the TRANSFERABLE KNOWLEDGE as an EXECUTABLE SOLUTION SKELETON \ +that would help a reader solve SIMILAR problems in the same domain. + +Your output is the ENTIRE useful content — there is no expansion tool, no second pass. \ +The reader will apply this knowledge to a DIFFERENT problem, so focus on what transfers. + +Principles: +1. OUTPUT AN EXECUTABLE STEP CHAIN: numbered steps that a solver can directly follow. \ +Each step should state WHAT to do and HOW (with the formula/technique), not just \ +name the concept. +2. INCLUDE FULL FORMULAS: theorems, identities, inequalities — state each \ +with its COMPLETE MATHEMATICAL EXPRESSION, not just the name. +3. STATE APPLICABILITY: what structural features of a problem signal that this \ +approach works (e.g. "when the constraint is a sum of squares"). +4. PRESERVE KEY INSIGHTS: the non-obvious ideas or tricks that make the approach \ +work — the things a solver would NOT think of without guidance. +5. REMOVE: problem-specific numeric calculations, dead-end explorations, \ +hesitations, verbose restatements, and trivial arithmetic. +6. FORMAT: Start with a one-line "Applicability" statement, then numbered steps, \ +then key formulas. Keep it concise and actionable. +7. NO meta-commentary about the compression process. NO preamble. +""" + +COMPRESS_USER = ( + '## Reader Problem (context only — do NOT solve it)\n{query}\n\n' + '## Reasoning Trace to Condense\n{text}') + + +def _is_truncated_compression(text: str) -> bool: + if not text or not text.strip(): + return True + lines = [l.strip() for l in text.strip().splitlines() if l.strip()] + if len(lines) < 3: + return True + last_line = lines[-1] + # Truncated if last line looks incomplete (no terminal punctuation/formula) + if last_line and last_line[-1] not in '.。!！)）]】}\\$': + # Allow lines ending with numbers, boxed answers, etc. + if not re.search(r'\d$|\\boxed|\$|\)$', last_line): + return True + return False + + +# -- API rate limiter ---------------------------------------------------------- +_api_semaphore = threading.Semaphore(CONDENSE_API_CONCURRENCY) +_api_bucket_lock = threading.Lock() +_api_tokens = [float(CONDENSE_API_CONCURRENCY)] +_api_last_refill = [time.monotonic()] + + +def _api_throttle(): + _api_semaphore.acquire() + wait = 0.0 + try: + with _api_bucket_lock: + now = time.monotonic() + elapsed = now - _api_last_refill[0] + refill = elapsed / CONDENSE_API_MIN_INTERVAL + _api_tokens[0] = min(float(CONDENSE_API_CONCURRENCY), _api_tokens[0] + refill) + _api_last_refill[0] = now + if _api_tokens[0] >= 1.0: + _api_tokens[0] -= 1.0 + else: + wait = (1.0 - _api_tokens[0]) * CONDENSE_API_MIN_INTERVAL + _api_tokens[0] = 0.0 + finally: + _api_semaphore.release() + if wait > 0: + time.sleep(wait) + + +def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Optional[str]: + _api_throttle() + trajectory = {'messages': messages} + sp = TwinkleSamplingParams(temperature=CONDENSE_TEMPERATURE, max_tokens=CONDENSE_MAX_TOKENS) + try: + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + except Exception as exc: + logger.warning(f'[condense-api] error: {exc}') + return None + content = (reply.get('content') or '').strip() + if not content: + return None + m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) + if m: + content = m.group(1).strip() + return content + + +def _api_hint_analysis_batch( + api_client: OpenAIClient, + problems: List[str], + condensed_examples: List[List[Dict[str, str]]], +) -> List[Optional[str]]: + """Call API to pre-analyze RAG relevance for each problem.""" + _MAX_HINT_INPUT = 8000 + results: List[Optional[str]] = [None] * len(problems) + tasks = [] + for i, prob in enumerate(problems): + if not condensed_examples[i]: + continue + traces = [ex.get('thinking', '') for ex in condensed_examples[i]] + merged_thinking = '\n---\n'.join(traces) + if len(merged_thinking) > _MAX_HINT_INPUT: + merged_thinking = merged_thinking[:_MAX_HINT_INPUT] + '\n[...truncated]' + user_msg = HINT_ANALYSIS_USER.format(query=prob, thinking=merged_thinking) + msgs = [ + {'role': 'system', 'content': HINT_ANALYSIS_SYSTEM}, + {'role': 'user', 'content': user_msg}, + ] + tasks.append((i, msgs)) + + if not tasks: + return results + + def _call_one(idx, msgs): + _api_throttle() + try: + trajectory = {'messages': msgs} + sp = TwinkleSamplingParams( + temperature=HINT_ANALYSIS_TEMPERATURE, + max_tokens=HINT_ANALYSIS_MAX_TOKENS) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + content = (reply.get('content') or '').strip() + # Treat "No applicable methods." as empty (will trigger fallback) + if not content or content == 'No applicable methods.': + return idx, None + return idx, content + except Exception as exc: + logger.warning(f'[hint-analysis] error for idx={idx}: {exc}') + return idx, None + + with ThreadPoolExecutor(max_workers=min(len(tasks), CONDENSE_API_CONCURRENCY)) as pool: + futs = [pool.submit(_call_one, idx, msgs) for idx, msgs in tasks] + for fut in as_completed(futs): + idx, analysis = fut.result() + results[idx] = analysis + + n_success = sum(1 for r in results if r) + logger.info(f'[hint-analysis] completed {n_success}/{len(tasks)} analyses') + return results + + +# --------------------------------------------------------------------------- +# LLM-based decontamination +# --------------------------------------------------------------------------- + +_DECONTAM_JUDGE_PROMPT = ( + 'We are building a RAG-augmented math training system. Problem A is the test ' + 'question; Problem B was retrieved from a knowledge base.\n' + 'Answer YES only if A and B are essentially the SAME specific problem — ' + 'i.e. solving B directly gives you A\'s answer (just different wording/notation/' + 'format/negation).\n' + 'Answer NO if they merely share the same method/topic but have different ' + 'specific values, equations, or geometric configurations — learning B\'s ' + 'approach still requires independent work to solve A.\n' + 'Problem A: {prob_a}\n' + 'Problem B: {prob_b}\n' + 'Answer only YES or NO.' +) + + +def _llm_judge_same_problem( + api_client: OpenAIClient, pairs: List[tuple], +) -> List[bool]: + """Batch LLM judge: are (problem_a, problem_b) the same problem? + + Returns list of bools (True = same problem = should filter). + """ + if not pairs or not api_client: + return [False] * len(pairs) + + results = [False] * len(pairs) + + def _judge_one(idx, pa, pb): + prompt = _DECONTAM_JUDGE_PROMPT.format(prob_a=pa, prob_b=pb) + msgs = [{'role': 'user', 'content': prompt}] + _api_throttle() + try: + trajectory = {'messages': msgs} + sp = TwinkleSamplingParams(temperature=0.1, max_tokens=8) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + answer = (reply.get('content') or '').strip().upper() + return idx, 'YES' in answer + except Exception: + return idx, False + + with ThreadPoolExecutor(max_workers=min(len(pairs), CONDENSE_API_CONCURRENCY)) as pool: + futs = [pool.submit(_judge_one, i, pa, pb) for i, (pa, pb) in enumerate(pairs)] + for fut in as_completed(futs): + idx, is_same = fut.result() + results[idx] = is_same + return results + + +def _llm_decontaminate( + api_client: OpenAIClient, + problems: List[str], + all_examples: List[List[Dict[str, str]]], +) -> List[List[Dict[str, str]]]: + """Apply LLM-based decontamination: remove retrievals judged as same problem.""" + judge_pairs = [] # (qi, ret_idx, prob_a, prob_b) + for qi, exs in enumerate(all_examples): + for ri, ex in enumerate(exs): + judge_pairs.append((qi, ri, problems[qi], ex.get('query', ''))) + + if not judge_pairs: + return all_examples + + pairs_input = [(pa, pb) for _, _, pa, pb in judge_pairs] + verdicts = _llm_judge_same_problem(api_client, pairs_input) + to_remove = set() + for vi, (qi, ri, _, _) in enumerate(judge_pairs): + if verdicts[vi]: + to_remove.add((qi, ri)) + + if to_remove: + logger.info(f'[decontam-llm] filtered {len(to_remove)} same-problem retrievals') + for qi in range(len(all_examples)): + all_examples[qi] = [ + ex for ri, ex in enumerate(all_examples[qi]) + if (qi, ri) not in to_remove + ] + return all_examples + + +def condense_traces( + examples_batch: List[List[Dict[str, str]]], + problems: List[str], + api_client: OpenAIClient, + condenser_sampler=None, + compress_params=None, + special_tokens: set = None, + max_output_len: int = 2000, + dp_size: int = 1, +) -> List[List[Dict[str, str]]]: + """Compress retrieved thinking traces with query-aware condenser. + + Primary: local vLLM condenser (if provided). + Fallback: API condenser. + Final fallback: raw trace truncated to max_output_len. + """ + result: List[List[Dict[str, str]]] = [] + # Flatten all (batch_idx, ex_idx, problem, example) for batch processing + tasks = [] + for bi, (exs, prob) in enumerate(zip(examples_batch, problems)): + for ei, ex in enumerate(exs): + tasks.append((bi, ei, prob, ex)) + + if not tasks: + return [[] for _ in examples_batch] + + # Build condense prompts (aligned with make_embedding_dataset.py hard path) + prompts = [] + for _, _, prob, ex in tasks: + user_msg = COMPRESS_USER.format(query=prob, text=ex['thinking']) + prompts.append([{'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_msg}]) + + # Phase 1: local vLLM condenser + condensed = [None] * len(tasks) + condense_sources = ['raw'] * len(tasks) + fallback_indices = [] + + if condenser_sampler is not None and compress_params is not None: + sampler_inputs = [{'messages': p} for p in prompts] + # The local vLLM sampler runs data-parallel across ``dp_size`` workers + # and requires at least one item per worker (it errors with + # "Batch too small for N workers" otherwise). Pad the batch up to a + # multiple of dp_size by repeating the last item, run, then keep only + # the first ``n_real`` responses and drop the padding. + n_real = len(sampler_inputs) + pad_size = 0 + if dp_size > 1 and n_real > 0 and n_real % dp_size != 0: + pad_size = dp_size - (n_real % dp_size) + sampler_inputs = sampler_inputs + [sampler_inputs[-1]] * pad_size + try: + responses = condenser_sampler.sample(sampler_inputs, compress_params) + except Exception as exc: + logger.warning(f'[condense] sampler error: {exc}') + responses = [None] * len(sampler_inputs) + if pad_size: + responses = responses[:n_real] + for ri, resp in enumerate(responses): + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + if special_tokens: + for tok in special_tokens: + text = text.replace(tok, '') + text = text.rstrip() + if text and not _is_truncated_compression(text): + condensed[ri] = text + condense_sources[ri] = 'local' + else: + fallback_indices.append(ri) + else: + fallback_indices = list(range(len(tasks))) + + # Phase 2: API fallback + if fallback_indices and api_client: + with ThreadPoolExecutor(max_workers=CONDENSE_API_CONCURRENCY) as pool: + futures = {} + for ri in fallback_indices: + futures[pool.submit(_api_condense_single, api_client, prompts[ri])] = ri + for fut in as_completed(futures): + ri = futures[fut] + api_result = fut.result() + if api_result and not _is_truncated_compression(api_result): + condensed[ri] = api_result + condense_sources[ri] = 'api' + + # Phase 3: assemble results (fallback to raw truncation) + result = [[] for _ in examples_batch] + for ti, (bi, ei, prob, ex) in enumerate(tasks): + compressed = condensed[ti] + raw_len = len(ex['thinking']) + sim_val = ex.get('_sim', 0.0) + if compressed: + result[bi].append({'query': ex['query'], + 'thinking': _strip_condenser_markers(compressed), + '_condense_source': condense_sources[ti], + '_raw_trace_len': raw_len, '_sim': sim_val}) + else: + result[bi].append({'query': ex['query'], + 'thinking': ex['thinking'][:max_output_len], + '_condense_source': 'raw', + '_raw_trace_len': raw_len, '_sim': sim_val}) + + n_ok = sum(1 for c in condensed if c) + logger.info(f'[condense] {n_ok}/{len(tasks)} compressed ok, ' + f'{len(tasks) - n_ok} fell back to raw truncation') + return result + + +def _strip_condenser_markers(text: str) -> str: + """Light cleanup of condenser output. + + Removes any residual markdown headers or meta-lines that don't carry + solution content. Keeps numbered steps and equations intact. + """ + # Remove legacy ## headers if condenser still emits them + if '## More' in text: + text = text.split('## More', 1)[0] + text = re.sub(r'^##\s*Summary\s*\n?', '', text, flags=re.MULTILINE) + text = re.sub(r'^Topic:\s*.*\n?', '', text, flags=re.MULTILINE) + # Remove meta-commentary lines + text = re.sub(r'^\s*\(Note:.*\)\s*$', '', text, flags=re.MULTILINE) + return text.strip() + + +# --------------------------------------------------------------------------- +# Boxed answer extraction +# --------------------------------------------------------------------------- +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + """Extract the last \\boxed{...} content, handling nested braces.""" + if not text: + return None + last_match = None + for m in _BOXED_RE.finditer(text): + start = m.end() + depth = 1 + i = start + while i < len(text) and depth > 0: + if text[i] == '{': + depth += 1 + elif text[i] == '}': + depth -= 1 + i += 1 + if depth == 0: + last_match = text[start:i - 1].strip() + return last_match + + +def normalize_answer(ans: str) -> str: + """Normalize a math answer string for comparison.""" + if not ans: + return '' + s = ans.strip() + # MCQ: extract bare letter from \textbf{(D)}, \text{(A)}, \mathbb{A}, (B), etc. + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.replace(' ', '') + s = s.replace(r'\,', '') + s = s.replace(r'\;', '') + s = s.replace(r'\!', '') + s = s.replace(r'\text', '') + s = s.replace(r'\mathrm', '') + s = s.replace(r'\displaystyle', '') + s = re.sub(r'\\(?:left|right)[.()\[\]|]', '', s) + s = s.replace(r'\dfrac', r'\frac') + s = s.replace(r'\tfrac', r'\frac') + s = s.strip('$').strip() + # Strip unit-like brace suffixes: {cm}, {m}, {kg}, etc. + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + # Normalize degree: ^\circ, ^{\circ}, ° → deg + s = re.sub(r'\^\\circ|\^\{\\circ\}|°', 'deg', s) + # Canonicalize \frac{a}{b} → (a)/(b) + def _frac_to_slash(m): + # Handle nested braces in numerator/denominator + text = m.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 # skip '{' + den_start = pos + depth = 1 + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + denom = text[den_start:pos - 1] + return f'({numer})/({denom})' + s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) + # Also handle bare a/b → (a)/(b) for consistent comparison + # Only simple integer/variable fractions: 17/5 → (17)/(5) + s = re.sub(r'(? bool: + """Try to evaluate both as floats; match if within 1e-9 relative tolerance.""" + try: + va = float(a.replace('(', '').replace(')', '')) + vb = float(b.replace('(', '').replace(')', '')) + return abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + except (ValueError, ZeroDivisionError): + pass + # Try evaluating simple fraction expressions like (17)/(5) + frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') + def _eval_frac(s): + m = frac_re.match(s) + if m: + try: + return float(m.group(1)) / float(m.group(2)) + except (ValueError, ZeroDivisionError): + pass + return None + va, vb = _eval_frac(a), _eval_frac(b) + if va is not None and vb is not None: + return abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + return False + + +# MCQ compound pattern: \text{(D) }49, \textbf{(C)}12, (B) 21, etc. +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$' +) + + +def _split_mcq(ans: str): + """Split an MCQ answer into (letter, value) components. + + Handles compound forms (``\\text{(D) }49``, ``(B) 21``) as well as a + bare letter (``D`` -> letter only) and a bare value (``21`` -> value only). + Returns ``(letter_or_None, value_or_None)``. + """ + s = ans.strip() + m = _MCQ_REF_RE.match(s) + if m: + letter = m.group(1) or m.group(3) + value = (m.group(2) or m.group(4) or '').strip() + return letter, (value or None) + # Bare single letter (with optional \text/\textbf/\mathbb wrapper or parens). + # \mathbb{A} appears as a dirty reference label for option A in some rows. + bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if bl: + return bl.group(1), None + return None, s or None + + +def answers_match(predicted: str, reference: str) -> bool: + """Check if two math answers are equivalent. + + Supports bidirectional MCQ matching: either side may be a bare option + letter, a bare value, or a compound ``(letter) value`` form. The answer is + considered correct if the letters match, or if the values match. + """ + if not predicted or not reference: + return False + norm_p = normalize_answer(predicted) + norm_r = normalize_answer(reference) + if norm_p == norm_r: + return True + if _try_numeric_equal(norm_p, norm_r): + return True + + # Symmetric MCQ matching: decompose both sides into (letter, value). + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + + # Match on the option letter (only meaningful if both sides carry a letter). + if p_letter and r_letter and p_letter == r_letter: + return True + + # If both sides are letter-only (a bare option letter with no value), the + # letters are the only signal; differing letters mean a mismatch. Do NOT + # fall through to value comparison, which would spuriously match dirty + # labels like \mathbb{A} vs \mathbb{B}. + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + + # Match on the value part (compare whichever value each side exposes; fall + # back to the raw normalized string when a side has no separate value). + p_val = normalize_answer(p_value) if p_value else norm_p + r_val = normalize_answer(r_value) if r_value else norm_r + if p_val and r_val: + if p_val == r_val or _try_numeric_equal(p_val, r_val): + return True + return False + + +# --------------------------------------------------------------------------- +# Dataset loading +# --------------------------------------------------------------------------- + +def load_aops(n: int, seed: int = 42) -> List[Dict[str, Any]]: + """Load AoPS boxed problems, sample n, extract reference answers.""" + from modelscope import MsDataset + ds = MsDataset.load(AOPS_DATASET_ID, split='train', + download_mode='reuse_dataset_if_exists') + boxed = [] + for row in ds: + if not row['metadata'].get('boxed'): + continue + ref = extract_boxed(row['solution']) + if not ref: + continue + boxed.append({ + 'problem': row['problem'], + 'solution': row['solution'], + 'reference_answer': ref, + 'tags': row.get('tags', []), + }) + sys.stderr.write(f'[aops] {len(boxed)} boxed problems with extractable answers\n') + rng = random.Random(seed) + rng.shuffle(boxed) + if n > 0 and n < len(boxed): + boxed = boxed[:n] + sys.stderr.write(f'[aops] sampled {n} problems\n') + return boxed + + +def load_math(n: int, seed: int = 42, split: str = 'test', + per_level: int = 0) -> List[Dict[str, Any]]: + """Load the MATH (Hendrycks) dataset from local extracted JSON files. + + Each problem's reference answer is the ``\\boxed{}`` content of its + ``solution`` (MATH solutions always end in a boxed answer). + + Sampling is *stratified by level* so every difficulty (Level 1-5) is + represented equally — required to measure how RAG gain varies with + difficulty. ``per_level`` (if >0) fixes the count per level; otherwise + ``n`` is split evenly across the 5 levels. When both are 0, all problems + are returned. The final list is shuffled with ``seed`` so index order is + stable/comparable across direct vs rag runs. + """ + import glob + root = os.path.join(MATH_DATA_DIR, split) + files = glob.glob(os.path.join(root, '*', '*.json')) + if not files: + raise FileNotFoundError( + f'[math] no problems found under {root!r}; set MATH_DATA_DIR or ' + f'extract MATH.zip there') + + by_level: Dict[str, List[Dict[str, Any]]] = {} + n_no_box = 0 + for fp in files: + try: + with open(fp, 'r', encoding='utf-8') as fin: + row = json.load(fin) + except Exception: + continue + ref = extract_boxed(row.get('solution', '')) + if not ref: + n_no_box += 1 + continue + level = row.get('level', 'Unknown') + by_level.setdefault(level, []).append({ + 'problem': row['problem'], + 'solution': row['solution'], + 'reference_answer': ref, + 'level': level, + 'type': row.get('type', ''), + }) + + total = sum(len(v) for v in by_level.values()) + sys.stderr.write( + f'[math] {total} problems with boxed answers across ' + f'{len(by_level)} levels (skipped {n_no_box} without boxed)\n') + + levels = sorted(by_level.keys()) + rng = random.Random(seed) + + # Decide how many per level. + if per_level <= 0 and n > 0: + per_level = max(1, n // max(1, len(levels))) + + sampled: List[Dict[str, Any]] = [] + for lv in levels: + pool = by_level[lv] + rng.shuffle(pool) + take = pool if per_level <= 0 else pool[:per_level] + sampled.extend(take) + sys.stderr.write(f'[math] {lv}: took {len(take)}/{len(pool)}\n') + + rng.shuffle(sampled) + sys.stderr.write(f'[math] total sampled: {len(sampled)}\n') + return sampled + + +# --------------------------------------------------------------------------- +# Prompt building +# --------------------------------------------------------------------------- + +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.' +) + +RAG_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.\n\n' + 'You will first see example problem-solving traces or skills. ' + 'Learn from the reasoning methodology demonstrated in these examples, ' + 'then thinking to solve the actual problem.' +) + +RAG_FOLLOWUP = ( + 'The above is a reference solution to a similar problem. ' + 'You may use any applicable techniques from it, or ignore it ' + 'if you find a better approach. ' + 'Solve the problem step by step and put your final answer in \\boxed{}.' +) + +HINT_FOLLOWUP = ( + 'The above are applicable solution approaches extracted from similar problems. ' + 'You may use any applicable techniques from them, or ignore them ' + 'if you find a better approach. ' + 'Solve the problem step by step and put your final answer in \\boxed{}.' +) + +# Reminder appended to the final user turn. Without this, the reasoning model +# can loop indefinitely on multiple-choice problems, oscillating between boxing +# the option letter and boxing the value (e.g. "I'll box B. I'll box 21. ...") +# and never terminating. Boxing BOTH the letter and value removes the ambiguity +# (the grader accepts either), so the model has no format decision to agonize over. +MCQ_INSTRUCTION = ( + '\n\nNote: If the problem is multiple-choice (it lists options such as ' + '(A), (B), (C), ...), put BOTH the option letter and its value in the box, ' + 'e.g. \\boxed{(B) 21}. Otherwise, box the value directly. Decide the answer ' + 'format once and do not deliberate over which form to box.' +) + + +def build_direct_prompt(problem: str) -> Dict[str, Any]: + return { + 'messages': [ + {'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem + MCQ_INSTRUCTION}, + ] + } + + +def build_hint_prompt(problem: str, hint_analysis: str) -> Dict[str, Any]: + """Build prompt with pre-analyzed hint in a multi-turn conversation. + + Mirrors ``build_rag_prompt``: the hint is presented as an assistant + "extracted approaches" turn (instead of being buried in the system + prompt), followed by a user instruction that provides a clear closing + directive to solve the problem and box the answer. Keeping the final + solve/box instruction in a dedicated user turn (rather than in the + system prompt) helps the reasoning model terminate cleanly. + """ + messages: List[Dict[str, str]] = [ + {'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}, + {'role': 'assistant', + 'content': ('Here are applicable solution approaches extracted from ' + f'similar problems:\n\n{hint_analysis}')}, + {'role': 'user', 'content': HINT_FOLLOWUP + MCQ_INSTRUCTION}, + ] + return {'messages': messages} + + +def build_rag_prompt(problem: str, + examples: List[Dict[str, str]]) -> Dict[str, Any]: + """Approach B: multi-turn assistant format. + + The trace is presented as an assistant "retrieval" turn, followed by + a user instruction that constrains the model to use methodology only. + """ + messages: List[Dict[str, str]] = [{'role': 'system', 'content': DIRECT_SYSTEM}] + messages.append({'role': 'user', 'content': problem}) + # Build trace content from retrieved examples + trace_parts = [] + for i, ex in enumerate(examples, 1): + trace_parts.append(f'[Retrieved Example {i}]\nProblem: {ex["query"]}\n' + f'Reasoning:\n{ex["thinking"]}') + trace_text = '\n\n'.join(trace_parts) + messages.append({'role': 'assistant', + 'content': f'I found relevant reasoning traces from the knowledge base!\n\n{trace_text}'}) + messages.append({'role': 'user', 'content': RAG_FOLLOWUP + MCQ_INSTRUCTION}) + return {'messages': messages} + + +# --------------------------------------------------------------------------- +# 13-gram Jaccard decontamination +# --------------------------------------------------------------------------- + +def _normalize_for_ngram(text: str) -> str: + """Normalize text for n-gram comparison: strip LaTeX markup, lowercase.""" + text = text.lower() + text = re.sub(r'\$+', '', text) + text = re.sub(r'\\[a-z]+\{([^}]*)\}', r'\1', text) + text = re.sub(r'\\[a-z]+', ' ', text) + text = re.sub(r'[{}\\^_$]', '', text) + text = re.sub(r'\s+', ' ', text).strip() + return text + + +def _ngram_jaccard(text_a: str, text_b: str, n: int = 13) -> float: + """13-gram character-level Jaccard similarity.""" + a = _normalize_for_ngram(text_a) + b = _normalize_for_ngram(text_b) + if len(a) < n or len(b) < n: + return 0.0 + grams_a = set(a[i:i + n] for i in range(len(a) - n + 1)) + grams_b = set(b[i:i + n] for i in range(len(b) - n + 1)) + if not grams_a or not grams_b: + return 0.0 + return len(grams_a & grams_b) / len(grams_a | grams_b) + + +# --------------------------------------------------------------------------- +# Embedding / RAG helpers +# --------------------------------------------------------------------------- + +def _wrap_anchor(text: str) -> List[Dict[str, str]]: + return [ + {'role': 'user', 'content': text}, + {'role': 'assistant', 'content': 'Match the correct response here.'}, + ] + + +def get_embeddings(model: TransformersModel, template: Qwen3_5Template, + texts: List[str], dp_size: int) -> np.ndarray: + if not texts: + return np.zeros((0,), dtype=np.float32) + n = len(texts) + pad_n = (-n) % dp_size + padded = list(texts) + [' '] * pad_n if pad_n else list(texts) + features = [] + for t in padded: + feat = template.encode({'messages': _wrap_anchor(t or ' ')}) + feat['labels'] = [1] + features.append(feat) + out = model.forward_only(inputs=features, task='embedding', return_logits=True) + emb = out['embeddings'] + if isinstance(emb, torch.Tensor): + emb = emb.detach().to(torch.float32).cpu().numpy() + emb = np.asarray(emb, dtype=np.float32) + return emb[:n] if pad_n else emb + + +def retrieve_examples(tbl, query_vecs: np.ndarray, top_k: int, + use_thinking_raw: bool, sim_threshold: float = 0.0, + problems: List[str] = None, + decontam_threshold: float = 0.0, + ) -> List[List[Dict[str, str]]]: + thinking_field = 'thinking_raw' if use_thinking_raw else 'cot_compressed' + fetch_limit = top_k + 50 if decontam_threshold > 0 else top_k + n_queries = len(query_vecs) + all_examples: List[List[Dict[str, str]]] = [None] * n_queries + decontam_skipped = 0 + _decontam_lock = threading.Lock() + + def _search_one(qi: int): + nonlocal decontam_skipped + vec = query_vecs[qi] + results = ( + tbl.search(vec.astype(np.float32).tolist()) + .metric('dot') + .limit(fetch_limit) + .select(['query_raw', thinking_field, '_distance']) + .to_list() + ) + problem_text = problems[qi] if problems else '' + examples = [] + local_skipped = 0 + for r in results: + if len(examples) >= top_k: + break + sim = 1.0 - r.get('_distance', 0.0) + if sim < sim_threshold: + continue + q = r.get('query_raw', '') + t = r.get(thinking_field, '') + if not t: + continue + if decontam_threshold > 0 and problem_text and q: + ng_sim = _ngram_jaccard(problem_text, q) + if ng_sim > decontam_threshold: + local_skipped += 1 + continue + examples.append({'query': q, 'thinking': t, '_sim': round(sim, 4), + '_raw_trace_len': len(t)}) + all_examples[qi] = examples + if local_skipped: + with _decontam_lock: + decontam_skipped += local_skipped + + with ThreadPoolExecutor(max_workers=min(n_queries, 16)) as pool: + list(pool.map(_search_one, range(n_queries))) + + if decontam_skipped > 0: + logger.info(f'[decontam] skipped {decontam_skipped} leaked retrievals ' + f'(13-gram Jaccard > {decontam_threshold})') + return all_examples + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--mode', choices=['direct', 'rag'], default='rag') + p.add_argument('--dataset', choices=['aops', 'math'], default='math', + help='Evaluation dataset. "math" = MATH (Hendrycks), ' + 'stratified by level for a difficulty-vs-gain curve.') + p.add_argument('--math-split', default='test', + help='MATH split to load (test/train).') + p.add_argument('--per-level', type=int, default=100, + help='MATH only: problems per difficulty level (default 100 ' + '-> 500 total across Level 1-5). If 0, --n is split ' + 'evenly across the 5 levels.') + p.add_argument('--n', type=int, default=0, + help='Pool size: sample this many problems (0 = all boxed). ' + 'In RAG mode with --target-eval, set this to 0 for max coverage.') + p.add_argument('--target-eval', type=int, default=0, + help='Stop after this many problems are successfully evaluated ' + '(0 = no limit, evaluate the entire sampled set — the ' + 'default, so all 500 stratified MATH problems are run). ' + 'RAG mode: counts problems with valid traces after ' + 'decontam; direct mode: ignored, evaluates all filtered.') + p.add_argument('--db-path', default='./output.oldemb/thinking_rag/lance.db') + p.add_argument('--table', default='thinking_traces') + p.add_argument('--top-k', type=int, default=1) + p.add_argument('--use-cot-compressed', action='store_true', + help='Use pre-compressed cot_compressed field instead of thinking_raw.') + p.add_argument('--sim-threshold', type=float, default=0.75, + help='Minimum cosine similarity for retrieved traces. ' + 'Traces below this are discarded at retrieval time.') + p.add_argument('--decontam-threshold', type=float, default=0.20, + help='13-gram Jaccard threshold for leak detection. ' + 'Retrieved traces above this are skipped (0=disabled).') + p.add_argument('--llm-decontam', action='store_true', default=True, + help='LLM-based decontamination (default ON): API judges whether ' + 'retrieved problem is the same as the test problem. ' + 'Applied after 13-gram decontam, before condensing. ' + 'Use --no-llm-decontam to disable.') + p.add_argument('--no-llm-decontam', dest='llm_decontam', action='store_false', + help='Disable LLM-based decontamination.') + p.add_argument('--max-trace-len', type=int, default=12000) + p.add_argument('--condense', action='store_true', default=True, + help='Enable condenser re-compression on retrieved traces ' + '(default ON). Use --no-condense to inject raw traces.') + p.add_argument('--no-condense', dest='condense', action='store_false', + help='Disable condenser; inject raw retrieved traces.') + p.add_argument('--condense-max-len', type=int, default=2000, + help='Max chars of condensed trace (fallback truncation).') + p.add_argument('--batch-size', type=int, default=16) + p.add_argument('--seed', type=int, default=42) + p.add_argument('--hint', action='store_true', default=False, + help='Enable API hint filtering on retrieved traces (default OFF; ' + 'raw RAG injects the condensed trace directly). ' + 'In rag mode: retrieve → condense → API filters trace → refined system prompt. ' + 'In direct mode: ignored (no traces to filter).') + p.add_argument('--no-hint', dest='hint', action='store_false', + help='Disable API hint filtering; inject condensed trace directly.') + p.add_argument('--problem-ids-file', default=None, + help='File listing problem indices evaluated by RAG mode. ' + 'RAG mode writes this file; direct mode reads it to ' + 'evaluate the same subset (use --no-filter to disable). ' + 'Defaults to a dataset-specific path.') + p.add_argument('--no-filter', action='store_true', + help='In direct mode, evaluate ALL sampled problems ' + 'instead of filtering to RAG subset.') + p.add_argument('--output', default=None) + args = p.parse_args() + + # Dataset-specific default paths (keeps aops and math runs from colliding). + if args.problem_ids_file is None: + args.problem_ids_file = ( + f'./output/thinking_rag/{args.dataset}_rag_problem_ids.json') + + if args.output is None: + suffix = f'{args.mode}_hint' if (args.hint and args.mode == 'rag') else args.mode + args.output = ( + f'./output/thinking_rag/{args.dataset}_{suffix}_results.jsonl') + + if args.condense and args.use_cot_compressed: + logger.warning('--condense requires thinking_raw, ignoring --use-cot-compressed') + args.use_cot_compressed = False + + if args.dataset == 'math': + records = load_math(n=args.n, seed=args.seed, split=args.math_split, + per_level=args.per_level) + else: + records = load_aops(n=args.n, seed=args.seed) + + is_rag = (args.mode == 'rag') + + # Direct mode: filter to same problems RAG evaluated (controlled comparison) + original_indices = list(range(len(records))) # track original indices + if not is_rag and not args.no_filter: + if os.path.exists(args.problem_ids_file): + with open(args.problem_ids_file) as f: + content = f.read().strip() + if content.startswith('['): + valid_indices = set(json.loads(content)) + else: + valid_indices = set(int(line) for line in content.splitlines() if line.strip()) + filtered = [(i, r) for i, r in enumerate(records) if i in valid_indices] + original_indices = [i for i, _ in filtered] + records = [r for _, r in filtered] + sys.stderr.write( + f'[direct] filtered to {len(records)} problems ' + f'from {args.problem_ids_file}\n') + else: + sys.stderr.write( + f'[direct] WARNING: {args.problem_ids_file} not found, ' + f'running all {len(records)} problems\n') + + condenser_gpus = int(os.environ.get('EVAL_CONDENSER_GPUS', 0)) if args.condense else 0 + + # Raw RAG relies on the API condenser (qwen3.7-max). Fail fast with a clear + # message if it's enabled without an API key and without a local condenser. + if is_rag and args.condense and not CONDENSE_API_KEY and condenser_gpus == 0: + sys.stderr.write( + '[condense] ERROR: --condense is ON but COMPRESS_API_KEY is unset ' + 'and no local condenser (EVAL_CONDENSER_GPUS=0).\n' + ' Fix one of:\n' + ' - export COMPRESS_API_KEY=sk-... (use API condenser)\n' + ' - EVAL_CONDENSER_GPUS=2 python ... (use local vLLM condenser)\n' + ' - pass --no-condense (inject raw traces)\n') + sys.exit(1) + + if is_rag: + num_gpus = EMB_GPUS + GEN_GPUS + condenser_gpus + device_groups = [ + DeviceGroup(name='emb_model', ranks=list(range(EMB_GPUS)), + device_type='GPU'), + DeviceGroup(name='sampler', + ranks=list(range(EMB_GPUS, EMB_GPUS + GEN_GPUS)), + device_type='GPU', gpus_per_worker=GEN_GPUS), + ] + if condenser_gpus > 0: + cond_start = EMB_GPUS + GEN_GPUS + device_groups.append( + DeviceGroup(name='condenser', + ranks=list(range(cond_start, cond_start + condenser_gpus)), + device_type='GPU')) + emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) + gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=num_gpus, + groups=device_groups, lazy_collect=False) + else: + device_groups = [ + DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), + device_type='GPU', gpus_per_worker=GEN_GPUS), + ] + gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=GEN_GPUS, + groups=device_groups, lazy_collect=False) + + sampler = vLLMSampler( + model_id=GEN_MODEL_ID, + engine_args={ + 'gpu_memory_utilization': GEN_GPU_MEM, + 'max_model_len': GEN_MAX_MODEL_LEN, + }, + device_mesh=gen_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=GEN_MAX_MODEL_LEN) + sys.stderr.write(f'[aops] vLLM sampler ready (model={GEN_MODEL_ID})\n') + + gen_params = TwinkleSamplingParams( + max_tokens=GEN_MAX_TOKENS, + temperature=GEN_TEMPERATURE, + top_p=GEN_TOP_P, + num_samples=1, + ) + + emb_model = emb_template = tbl = None + if is_rag: + import lancedb + db = lancedb.connect(args.db_path) + if args.table not in db.table_names(): + raise SystemExit(f'Table "{args.table}" not found in {args.db_path}') + tbl = db.open_table(args.table) + sys.stderr.write(f'[aops] LanceDB rows={tbl.count_rows()}\n') + + emb_model = TransformersModel( + model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, + remote_group='emb_model') + emb_model.set_processor(InputProcessor) + emb_model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) + emb_template = Qwen3_5Template( + model_id=EMBED_MODEL_ID, max_length=EMBED_MAX_LENGTH, + truncation_strategy='delete', enable_thinking=False) + sys.stderr.write('[aops] embedding model ready\n') + + # -- Condenser setup (API primary + optional local vLLM) ------------------- + condenser_api_client = None + condenser_sampler_obj = None + condenser_params = None + condenser_special_tokens = None + + if args.condense and is_rag: + condenser_api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + sys.stderr.write(f'[condense] API client ready (model={CONDENSE_API_MODEL})\n') + + if condenser_gpus > 0: + condenser_mesh = DeviceMesh.from_sizes( + world_size=condenser_gpus, dp_size=condenser_gpus) + condenser_sampler_obj = vLLMSampler( + model_id=CONDENSE_MODEL_ID, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': 32768}, + device_mesh=condenser_mesh, + remote_group='condenser', + ) + condenser_sampler_obj.set_template( + 'Qwen3_5Template', model_id=CONDENSE_MODEL_ID, + enable_thinking=False, truncation_strategy='delete', + max_length=32768) + condenser_template = Qwen3_5Template( + model_id=CONDENSE_MODEL_ID, max_length=32768, + enable_thinking=False, truncation_strategy='delete') + condenser_special_tokens = set(condenser_template.tokenizer.all_special_tokens) + condenser_params = TwinkleSamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, + temperature=CONDENSE_TEMPERATURE, + top_p=0.5, num_samples=1) + sys.stderr.write(f'[condense] local vLLM ready (model={CONDENSE_MODEL_ID})\n') + + # -- Hint analysis API client (reuses condenser API config) ----------------- + hint_api_client = None + if args.hint and is_rag: + if condenser_api_client is not None: + hint_api_client = condenser_api_client + else: + hint_api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + sys.stderr.write(f'[hint] API hint analysis enabled (model={CONDENSE_API_MODEL})\n') + + # -- LLM decontam API client --------------------------------------------------- + decontam_api_client = None + if args.llm_decontam and is_rag: + if hint_api_client is not None: + decontam_api_client = hint_api_client + elif condenser_api_client is not None: + decontam_api_client = condenser_api_client + else: + decontam_api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + sys.stderr.write(f'[decontam-llm] LLM decontamination enabled (model={CONDENSE_API_MODEL})\n') + + correct_count = 0 + total_count = 0 + skipped_indices: List[int] = [] # problems skipped by RAG (no valid trace) + evaluated_indices: List[int] = [] # problems actually evaluated + debug_records: List[Dict[str, Any]] = [] + + os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) + out_f = open(args.output, 'w', encoding='utf-8') + + # Open problem-ids files for incremental writing (RAG mode only) + ids_f = None + skip_f = None + if is_rag: + os.makedirs(os.path.dirname(args.problem_ids_file) or '.', exist_ok=True) + ids_f = open(args.problem_ids_file, 'w', encoding='utf-8') + skip_path = args.problem_ids_file.replace('.json', '_skipped.json') + skip_f = open(skip_path, 'w', encoding='utf-8') + + # -- RAG batch preparation (embed + retrieve + decontam + condense + hint) -- + def _prepare_rag_batch(batch_start: int): + """Prepare a RAG batch: returns (prompts, batch, all_examples, + hint_analyses, kept_global_indices, batch_skipped_indices) or None.""" + batch_end = min(batch_start + args.batch_size, len(records)) + batch = records[batch_start:batch_end] + problems = [r['problem'] for r in batch] + + query_vecs = get_embeddings(emb_model, emb_template, problems, EMB_GPUS) + use_raw = not args.use_cot_compressed + all_examples = retrieve_examples(tbl, query_vecs, args.top_k, + use_raw, args.sim_threshold, + problems=problems, + decontam_threshold=args.decontam_threshold) + if args.use_cot_compressed: + for exs in all_examples: + for ex in exs: + ex['thinking'] = _strip_condenser_markers(ex['thinking']) + + if args.llm_decontam and decontam_api_client: + all_examples = _llm_decontaminate( + decontam_api_client, problems, all_examples) + + if args.condense and condenser_api_client: + all_examples = condense_traces( + all_examples, problems, condenser_api_client, + condenser_sampler=condenser_sampler_obj, + compress_params=condenser_params, + special_tokens=condenser_special_tokens, + max_output_len=args.condense_max_len, + dp_size=condenser_gpus) + + hint_analyses = None + if args.hint and hint_api_client: + hint_analyses = _api_hint_analysis_batch( + hint_api_client, problems, all_examples) + + keep_mask = [] + for pi, (r, examples) in enumerate(zip(batch, all_examples)): + if not examples: + keep_mask.append(False) + elif hint_analyses and hint_analyses[pi]: + keep_mask.append(True) + else: + usable = [ex for ex in examples + if len(ex['thinking']) <= args.max_trace_len] + keep_mask.append(bool(usable)) + + batch_skipped = [] + for pi, keep in enumerate(keep_mask): + if not keep: + batch_skipped.append(batch_start + pi) + + kept_batch = [] + kept_examples = [] + kept_hints = [] + kept_global_indices = [] + for pi, keep in enumerate(keep_mask): + if keep: + kept_batch.append(batch[pi]) + kept_examples.append(all_examples[pi]) + kept_hints.append(hint_analyses[pi] if hint_analyses else None) + kept_global_indices.append(batch_start + pi) + + if not kept_batch: + return None, None, None, None, None, batch_skipped + + prompts = [] + for pi, (r, examples) in enumerate(zip(kept_batch, kept_examples)): + if kept_hints[pi]: + prompts.append(build_hint_prompt(r['problem'], kept_hints[pi])) + else: + filtered = [{'query': ex['query'], 'thinking': ex['thinking']} + for ex in examples + if len(ex['thinking']) <= args.max_trace_len] + prompts.append(build_rag_prompt(r['problem'], filtered)) + + return prompts, kept_batch, kept_examples, kept_hints, kept_global_indices, batch_skipped + + target_reached = False + batch_starts = list(range(0, len(records), args.batch_size)) + + if is_rag: + # Pipeline: prefetch next batch while current batch generates + from concurrent.futures import Future + prefetch_pool = ThreadPoolExecutor(max_workers=1) + # Prepare first batch synchronously + cur_result = _prepare_rag_batch(batch_starts[0]) + + for bi, batch_start in enumerate(batch_starts): + if target_reached: + break + prompts, batch, all_examples, hint_analyses, kept_global_indices, batch_skipped = cur_result + skipped_indices.extend(batch_skipped or []) + if skip_f and batch_skipped: + for sid in batch_skipped: + skip_f.write(f'{sid}\n') + skip_f.flush() + + # Submit next batch preparation in background + next_future: Optional[Future] = None + if bi + 1 < len(batch_starts) and not target_reached: + next_future = prefetch_pool.submit(_prepare_rag_batch, batch_starts[bi + 1]) + + if prompts is None: + # Entire batch skipped + cur_result = next_future.result() if next_future else None + continue + + # Generate (runs on gen GPU while next batch prepares on emb GPU + API) + responses = sampler.sample(prompts, gen_params) + + for i, (rec, resp) in enumerate(zip(batch, responses)): + seq = resp.sequences[0] if resp and resp.sequences else None + raw_output = '' + if seq is not None: + raw_output = seq.decoded or '' + raw_output = re.sub(r'<\|[^|]+\|>', '', raw_output).rstrip() + + predicted = extract_boxed(raw_output) + is_correct = answers_match(predicted, rec['reference_answer']) + if is_correct: + correct_count += 1 + total_count += 1 + + global_idx = kept_global_indices[i] + evaluated_indices.append(global_idx) + if ids_f: + ids_f.write(f'{global_idx}\n') + ids_f.flush() + + debug_rec = { + 'idx': global_idx, + 'reference_answer': rec['reference_answer'], + 'predicted': predicted, + 'is_correct': is_correct, + 'problem': rec['problem'], + 'model_output': raw_output, + } + if rec.get('level'): + debug_rec['level'] = rec['level'] + if rec.get('type'): + debug_rec['type'] = rec['type'] + debug_rec['num_traces'] = len(all_examples[i]) + if all_examples[i]: + ex0 = all_examples[i][0] + debug_rec['similarity'] = ex0.get('_sim', 0.0) + debug_rec['retrieved_query'] = ex0.get('query', '') + debug_rec['raw_trace_len'] = ex0.get('_raw_trace_len', 0) + debug_rec['condensed_trace'] = ex0['thinking'] + debug_rec['condensed_trace_len'] = len(ex0['thinking']) + debug_rec['condense_source'] = ex0.get('_condense_source', '') + if hint_analyses and hint_analyses[i]: + debug_rec['hint_analysis'] = hint_analyses[i] + debug_records.append(debug_rec) + out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') + out_f.flush() + + acc = correct_count / total_count if total_count else 0 + sys.stderr.write( + f' [{total_count}/{args.target_eval}] ' + f'acc={acc:.4f} ({correct_count}/{total_count})\n') + + if args.target_eval > 0 and total_count >= args.target_eval: + target_reached = True + + # Collect prefetched result for next iteration (skip if done) + if not target_reached and next_future: + cur_result = next_future.result() + else: + cur_result = None + + prefetch_pool.shutdown(wait=True) + else: + # Direct mode: no pipeline needed, just batch generate + for batch_start in batch_starts: + batch_end = min(batch_start + args.batch_size, len(records)) + batch = records[batch_start:batch_end] + prompts = [build_direct_prompt(r['problem']) for r in batch] + + responses = sampler.sample(prompts, gen_params) + + for i, (rec, resp) in enumerate(zip(batch, responses)): + seq = resp.sequences[0] if resp and resp.sequences else None + raw_output = '' + if seq is not None: + raw_output = seq.decoded or '' + raw_output = re.sub(r'<\|[^|]+\|>', '', raw_output).rstrip() + + predicted = extract_boxed(raw_output) + is_correct = answers_match(predicted, rec['reference_answer']) + if is_correct: + correct_count += 1 + total_count += 1 + + global_idx = original_indices[batch_start + i] + evaluated_indices.append(global_idx) + + debug_rec = { + 'idx': global_idx, + 'reference_answer': rec['reference_answer'], + 'predicted': predicted, + 'is_correct': is_correct, + 'problem': rec['problem'], + 'model_output': raw_output, + } + if rec.get('level'): + debug_rec['level'] = rec['level'] + if rec.get('type'): + debug_rec['type'] = rec['type'] + debug_records.append(debug_rec) + out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') + out_f.flush() + + acc = correct_count / total_count if total_count else 0 + sys.stderr.write( + f' [{total_count}/{len(records)}] ' + f'acc={acc:.4f} ({correct_count}/{total_count})\n') + + overall_acc = correct_count / total_count if total_count else 0 + print(f'\n{"=" * 60}') + print(f'{args.dataset.upper()} — mode={args.mode}, model={GEN_MODEL_ID}') + print(f' n={total_count}, seed={args.seed}') + if is_rag: + print(f' evaluated={len(evaluated_indices)}, skipped={len(skipped_indices)}') + print(f'{"=" * 60}') + print(f'Overall accuracy: {overall_acc:.4f} ({correct_count}/{total_count})') + + # Per-level breakdown (MATH: the difficulty-vs-gain curve we care about). + if any(r.get('level') for r in debug_records): + from collections import defaultdict + per = defaultdict(lambda: [0, 0]) # level -> [correct, total] + for r in debug_records: + lv = r.get('level', 'Unknown') + per[lv][1] += 1 + if r['is_correct']: + per[lv][0] += 1 + print(f'\nPer-level accuracy:') + for lv in sorted(per.keys()): + c, t = per[lv] + print(f' {lv:>10}: {c/t:.4f} ({c}/{t})') + + out_f.close() + print(f'\n[output] {len(debug_records)} records saved to {args.output}') + + if ids_f: + ids_f.close() + print(f'[output] problem IDs ({len(evaluated_indices)}) saved to {args.problem_ids_file}') + if skip_f: + skip_f.close() + if skipped_indices: + print(f'[output] skipped IDs ({len(skipped_indices)}) saved to ' + f'{args.problem_ids_file.replace(".json", "_skipped.json")}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/eval_math_by_level.sh b/cookbook/exp/embedding/eval_math_by_level.sh new file mode 100755 index 00000000..d4f8bbdb --- /dev/null +++ b/cookbook/exp/embedding/eval_math_by_level.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# MATH (Hendrycks) difficulty-stratified evaluation. +# +# Goal: measure how the (raw) RAG gain over direct varies with problem +# difficulty (Level 1-5). Runs raw RAG first (retrieve -> qwen3.7-max condense +# -> inject, no hint filtering; it writes the problem-id file), then direct on +# the *same* problems for a paired comparison. +# +# Usage: +# COMPRESS_API_KEY=sk-xxx bash cookbook/exp/embedding/eval_math_by_level.sh +# +# Env knobs: +# PER_LEVEL problems per difficulty level (default 100 -> 500 total) +# SEED stratified-sampling seed (default 100; must match across runs) +# DB_PATH LanceDB retrieval index +# SIM / TOPK retrieval threshold / top-k + +set -euo pipefail + +export COMPRESS_API_KEY="${COMPRESS_API_KEY:?Set COMPRESS_API_KEY}" + +SCRIPT="cookbook/exp/embedding/eval_gpqa_rag.py" +PER_LEVEL="${PER_LEVEL:-100}" +SEED="${SEED:-100}" +SIM="${SIM:-0.75}" +TOPK="${TOPK:-1}" +OUTDIR="./output/thinking_rag" +DB_PATH="${DB_PATH:-./output.oldemb/thinking_rag/lance.db}" + +mkdir -p "$OUTDIR" + +echo "============================================================" +echo " MATH by level: raw RAG (qwen3.7-max condenser, no hint)" +echo " per_level=$PER_LEVEL seed=$SEED" +echo "============================================================" +python "$SCRIPT" \ + --dataset math --math-split test \ + --mode rag \ + --per-level "$PER_LEVEL" --seed "$SEED" \ + --db-path "$DB_PATH" \ + --sim-threshold "$SIM" --top-k "$TOPK" \ + --condense \ + --output "$OUTDIR/math_rag_results.jsonl" + +echo "" +echo "============================================================" +echo " MATH by level: Direct (same problems as raw RAG)" +echo "============================================================" +# Direct reads math_rag_problem_ids.json (written above) to match the subset. +python "$SCRIPT" \ + --dataset math --math-split test \ + --mode direct \ + --per-level "$PER_LEVEL" --seed "$SEED" \ + --output "$OUTDIR/math_direct_results.jsonl" + +echo "" +echo "============================================================" +echo " Done. Compare with: python cookbook/exp/embedding/compare_math_levels.py" +echo "============================================================" diff --git a/cookbook/exp/embedding/eval_rag_recall.py b/cookbook/exp/embedding/eval_rag_recall.py new file mode 100644 index 00000000..19bc5ad6 --- /dev/null +++ b/cookbook/exp/embedding/eval_rag_recall.py @@ -0,0 +1,187 @@ +"""Self-recall evaluation: sample rows from LanceDB, re-encode query, check retrieval. + +Unlike the full build pipeline (which needs 8 GPUs for condenser + embedding), +this script only needs the embedding model (4 GPUs) since it uses the +already-compressed ``query_compressed`` stored in the index. + +Launch: + python cookbook/exp/embedding/eval_rag_recall.py + python cookbook/exp/embedding/eval_rag_recall.py --n 200 --top-k 20 + python cookbook/exp/embedding/eval_rag_recall.py --db-path ./output/thinking_rag/lance.db +""" +import argparse +import json +import os +import random +import sys +from typing import Any, Dict, List, Tuple + +import numpy as np +import torch + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.loss import InfonceLoss +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.template import Qwen3_5Template + +logger = get_logger() + +EMBED_MODEL_ID = os.environ.get( + 'EMBED_MODEL_ID', 'output/embedding_full_transformers/last-checkpoint') +EMB_GPUS = int(os.environ.get('EMB_GPUS', 4)) +EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 8192)) + + +def _wrap_anchor(text: str) -> List[Dict[str, str]]: + return [ + {'role': 'user', 'content': text}, + {'role': 'assistant', 'content': 'Match the correct response here.'}, + ] + + +def get_embeddings(model: TransformersModel, template: Qwen3_5Template, + texts: List[str]) -> np.ndarray: + if not texts: + return np.zeros((0,), dtype=np.float32) + n = len(texts) + pad_n = (-n) % EMB_GPUS + padded = list(texts) + [' '] * pad_n if pad_n else list(texts) + features = [] + for t in padded: + feat = template.encode({'messages': _wrap_anchor(t or ' ')}) + feat['labels'] = [1] + features.append(feat) + out = model.forward_only(inputs=features, task='embedding', return_logits=True) + emb = out['embeddings'] + if isinstance(emb, torch.Tensor): + emb = emb.detach().to(torch.float32).cpu().numpy() + emb = np.asarray(emb, dtype=np.float32) + return emb[:n] if pad_n else emb + + +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--db-path', default='./output/thinking_rag/lance.db') + p.add_argument('--table', default='thinking_traces') + p.add_argument('--n', type=int, default=100, help='Number of samples to probe.') + p.add_argument('--top-k', type=int, default=10) + p.add_argument('--seed', type=int, default=42) + p.add_argument('--batch-size', type=int, default=32) + p.add_argument('--output', default='./output/thinking_rag/recall_debug.jsonl', + help='JSONL file to dump per-sample debug info.') + args = p.parse_args() + + import lancedb + db = lancedb.connect(args.db_path) + if args.table not in db.table_names(): + raise SystemExit(f'Table "{args.table}" not found in {args.db_path}') + tbl = db.open_table(args.table) + total_rows = tbl.count_rows() + sys.stderr.write(f'[eval] table={args.table} rows={total_rows}\n') + + df = tbl.to_pandas() + n_sample = min(args.n, len(df)) + random.seed(args.seed) + sample_indices = random.sample(range(len(df)), n_sample) + samples = df.iloc[sample_indices].reset_index(drop=True) + sys.stderr.write(f'[eval] sampled {n_sample} rows for self-recall test\n') + + # Init embedding model only (no condenser needed). + device_groups = [ + DeviceGroup(name='emb_model', ranks=list(range(EMB_GPUS)), device_type='GPU'), + ] + emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=EMB_GPUS, groups=device_groups, + lazy_collect=False) + + model = TransformersModel(model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, + remote_group='emb_model') + model.set_processor(InputProcessor) + model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) + template = Qwen3_5Template(model_id=EMBED_MODEL_ID, max_length=EMBED_MAX_LENGTH, + truncation_strategy='delete', enable_thinking=False) + sys.stderr.write('[eval] embedding model ready\n') + + ks = sorted({1, 5, 10, args.top_k}) + hits = {k: 0 for k in ks} + per_source_hits: Dict[str, Dict[int, int]] = {} + per_source_total: Dict[str, int] = {} + debug_records: List[Dict[str, Any]] = [] + + # Batch encode and search. + for batch_start in range(0, n_sample, args.batch_size): + batch_end = min(batch_start + args.batch_size, n_sample) + batch = samples.iloc[batch_start:batch_end] + queries = batch['query_compressed'].tolist() + ids = batch['id'].tolist() + sources = batch['source'].tolist() + thinkings = batch['thinking_raw'].tolist() + query_raws = batch['query_raw'].tolist() + cot_compresseds = batch['cot_compressed'].tolist() + + anchor_emb = get_embeddings(model, template, queries) + + for i, (rid, src, vec) in enumerate(zip(ids, sources, anchor_emb)): + res = ( + tbl.search(vec.astype(np.float32).tolist()) + .metric('dot') + .limit(max(ks)) + .select(['id', 'source', 'query_compressed', 'cot_compressed', + 'thinking_raw', 'query_raw']) + .to_list() + ) + hit_ids = [item['id'] for item in res] + try: + rank = hit_ids.index(rid) + except ValueError: + rank = -1 + + for k in ks: + if 0 <= rank < k: + hits[k] += 1 + per_source_hits.setdefault(src, {kk: 0 for kk in ks})[k] += 1 + per_source_total[src] = per_source_total.get(src, 0) + 1 + per_source_hits.setdefault(src, {kk: 0 for kk in ks}) + + top1 = res[0] if res else {} + debug_records.append({ + 'id': rid, + 'source': src, + 'rank': rank, + 'query_raw': query_raws[i], + 'query_compressed': queries[i], + 'cot_compressed': cot_compresseds[i], + 'thinking_raw': thinkings[i][:2000], + 'top1_id': top1.get('id'), + 'top1_source': top1.get('source'), + 'top1_query_compressed': top1.get('query_compressed'), + 'top1_cot_compressed': top1.get('cot_compressed'), + 'top1_query_raw': top1.get('query_raw'), + 'top1_thinking_raw': (top1.get('thinking_raw') or '')[:2000], + 'top1_is_self': top1.get('id') == rid, + }) + + sys.stderr.write(f' probed {batch_end}/{n_sample}\n') + + print(f'\n=== Self-Recall @ k (n={n_sample}, seed={args.seed}) ===') + for k in ks: + print(f' recall@{k:<3} = {hits[k]/n_sample:.4f} ({hits[k]}/{n_sample})') + + print(f'\n=== Per-source recall@{max(ks)} ===') + for src in sorted(per_source_total, key=lambda s: -per_source_total[s]): + tot = per_source_total[src] + h = per_source_hits.get(src, {}).get(max(ks), 0) + print(f' {src:<48s} {h/tot:.4f} ({h}/{tot})') + + os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) + with open(args.output, 'w', encoding='utf-8') as f: + for rec in debug_records: + f.write(json.dumps(rec, ensure_ascii=False) + '\n') + print(f'\n[debug] {len(debug_records)} records saved to {args.output}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/make_embedding_dataset.py b/cookbook/exp/embedding/make_embedding_dataset.py new file mode 100644 index 00000000..847f222f --- /dev/null +++ b/cookbook/exp/embedding/make_embedding_dataset.py @@ -0,0 +1,758 @@ +"""Offline compression pipeline: raw datasets → condenser → pre-compressed embedding dataset. + +Loads think/index/hard datasets, compresses query/cot/negatives via vLLM condenser +with API fallback, saves a single HF Dataset ready for embedding training. + +Output schema: {anchor_text, positive_text, negative_texts, source} + +Launch (8 GPUs — 4 for vLLM condenser): + python cookbook/exp/embedding/make_embedding_dataset.py +""" +import json +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import Any, Dict, List, Optional + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle.utils.parallel import PosixFileLock +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from dataset_think import get_dataset as get_dataset_think # noqa: E402 +from dataset_index import get_dataset as get_dataset_index # noqa: E402 +from dataset_hard import get_dataset as get_dataset_hard # noqa: E402 + +logger = get_logger() + +# -- Model config ------------------------------------------------------------- +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +TEMPLATE_NAME = 'Qwen3_5Template' + +# -- GPU placement (condenser only) ------------------------------------------- +CONDENSER_GPUS = int(os.environ.get('CONDENSER_GPUS', 8)) + +# -- Dataset caps ------------------------------------------------------------- +TOTAL_SAMPLES: Optional[int] = None +THINK_CAP: Optional[int] = int(os.environ.get('THINK_CAP', 100_000)) +INDEX_CAP: Optional[int] = int(os.environ.get('INDEX_CAP', 100_000)) +HARD_CAP: Optional[int] = int(os.environ.get('HARD_CAP', 0)) or None +HARD_MAX_NEGATIVES = int(os.environ.get('HARD_MAX_NEGATIVES', 8)) + +# -- Compression params ------------------------------------------------------- +MIN_TEXT_CHARS = 256 +DATASET_MAX_TOKENS = 32768 +COMPRESS_TEMPERATURE = 0.2 +COMPRESS_TOP_P = 0.5 +COMPRESS_MAX_MODEL_LEN = 32768 +BATCH_SIZE = int(os.environ.get('COMPRESS_BATCH_SIZE', 128)) + +# -- API fallback ------------------------------------------------------------- +COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '') +COMPRESS_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') +COMPRESS_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') +API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) +API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 24)) +SAMPLER_TIMEOUT = float(os.environ.get('SAMPLER_TIMEOUT', 300)) + +# -- Output ------------------------------------------------------------------- +OUTPUT_DIR = os.environ.get('EMB_DATASET_OUTPUT', './output/embedding_dataset') +RESULTS_JSONL = f'{OUTPUT_DIR}/results.jsonl' +PROGRESS_FILE = f'{OUTPUT_DIR}/progress.json' + +# ============================================================================= +# Prompts +# ============================================================================= + +COMPRESS_SYSTEM = """\ +You are a compression and summary assistant. For the (query, source) pair, emit a Markdown \ +answer with TWO sections, designed to pair with the `extract_compressed` tool: \ +the reader absorbs `## Summary` directly, then calls `extract_compressed` \ +on any topic-key listed under `## More` to recover its \ +fuller content. + + `## Summary` — extreme-density text the reader reads directly. + `## More` — a topic index whose keys are valid arguments \ +to `extract_compressed` for recovering material not captured inline. + +Together the two sections must form a COMPLETE, NON-DISTORTING inventory of the \ +source for the query — nothing essential lost, nothing implied that the source \ +does not support. NO preamble, NO meta-commentary, NO code fences wrapping the \ +whole output. + +Output skeleton: + +## Summary +Topic: + + +## More +- : +- ... + +Format selection for the inline body (pick the MOST COMPACT form per query, mix \ +when helpful): +- Interface / signature → code notation directly: `func(a:int)->str` +- Factual / entity → telegraphic prose; drop function words; ":" for "is", "," \ +for "has" +- Skill / how-to / usage → lead with `Use when: `; numbered telegraphic \ +steps `1.do X 2.then Y`; close with `Output: ` when relevant +- Procedural → numbered short steps +- Analytical / design → hierarchical bullets with abbreviations + +`## Summary` rules: +1. TOPIC LINE — line 1 is ALWAYS `Topic: `, even when the \ +query is narrow. Anchors both the reader and the tool. +2. DENSITY — every token in the body carries query-relevant signal; cut filler. +3. PRIMARY-COMPLETE — never silently drop a fact essential to answering the \ +query. Anything cut for length MUST appear as a key under \ +`## More`. +4. NON-MISLEADING — phrasing must not let the reader infer anything the source \ +does not support; partial truths that mislead are worse than honest omissions \ +flagged in the index. +5. SELF-CONTAINED — the reader can act on the answer without re-opening the source. +6. FAITHFUL — only content the source supports; no fabrication, no extrapolation. +7. LANGUAGE — match the source language. +8. NO outer code fences around the whole answer; no meta-commentary. + +`## More` rules (MANDATORY — this section is never omitted): +1. FORMAT — each bullet is `- : `: + • topic-key — short, unambiguous, grounded in source vocabulary so the \ +`extract_compressed` tool can locate the aspect (e.g. `decorators`, \ +`error handling`, `pitfalls`). + • hint — tells WHAT the reader gains by expanding (concrete numbers, code \ +listings, secondary cases, edge details, related context, …); do NOT restate \ +the inline answer. +2. CRITERION — each bullet names an aspect that EXISTS in the source but is \ +NOT fully captured inline. Material that genuinely fits inline without \ +distortion MUST NOT be duplicated here. +3. FAITHFUL — hints must be grounded in the source; never speculate or invent. +4. ORDER — by relevance to the query, then by importance. +5. EMPTY CASE — if the source is so short / single-purpose that everything \ +fits inline, write a single line `- (none)`. + +Now begin.\ +""" + +COMPRESS_USER = ( + 'Downstream model will read your compressed block to decide whether to ' + 'expand it. Compress faithfully: preserve the passage topic + core facts. ' + 'Do NOT invent facts. Do NOT drop major facts. Do NOT write meta-commentary ' + 'about the Query (never write "Query info: absent", "no X mention", etc.); ' + 'if the passage does not address the Query, still summarize the passage. ' + 'CRITICAL LANGUAGE RULE: detect the dominant language of the Passage ' + '(NOT the Query, NOT this instruction) and write the ENTIRE output in that ' + 'same language; English passage → English output, Chinese passage → ' + 'Chinese output, Japanese passage → Japanese output. NEVER translate, ' + 'NEVER mix languages, NEVER copy these instructions into the output.\n\n' + '## Query (ordering hint only — still summarize the whole passage)\n{query}\n\n' + '## Passage\n{text}') + +EMBED_QUERY_Q = ( + 'Summarize this query for retrieval. ' + 'The body of ## Summary MUST follow this EXACT 4-line template — ' + 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the specific pattern, never generic labels.') + +EMBED_QUERY_COT = ( + 'Summarize this reasoning trace for retrieval. ' + 'The body of ## Summary MUST follow this EXACT 4-line template — ' + 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the specific pattern, never generic labels.') + +EMBED_QUERY_Q_LEGACY = ( + 'What problem does this passage address, and what skill or method is needed? ' + 'Topic must name the specific pattern, never generic labels. ' + 'Compress into a retrieval-friendly need description.') + +EMBED_QUERY_COT_LEGACY = ( + 'Extract the reusable skill: trigger conditions, key steps, and expected output. ' + 'Topic names the method/pattern; format as "Use when: ...", numbered steps, ' + '"Output: ...". Compress into a standardized procedure for retrieval.') + +EMBED_QUERY_REASONIR_Q = ( + 'Extract the abstract PROBLEM TYPE from this query. ' + 'IGNORE all specific numbers, values, variable names, and parameters — ' + 'focus ONLY on the CLASS of problem and the METHODOLOGY required. ' + 'The body of ## Summary MUST follow this EXACT 4-line template:\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the method class, never mention specific numbers.') + +EMBED_QUERY_REASONIR_COT = ( + 'Extract the abstract METHODOLOGY demonstrated in this solution. ' + 'IGNORE all specific numbers, values, and computed results — ' + 'focus ONLY on the general TECHNIQUE and key reasoning STEPS. ' + 'The body of ## Summary MUST follow this EXACT 4-line template:\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the method class, never mention specific numbers.') + + +# ============================================================================= +# Validation & API fallback +# ============================================================================= + +_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') +_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') + + +def _is_truncated_compression(text: str, schema: str = 'new') -> bool: + if not text or not text.strip(): + return True + if '## More' not in text or '## Summary' not in text: + return True + after_more = text.split('## More', 1)[1].strip() + if not after_more: + return True + last_line = after_more.splitlines()[-1].strip() + if not (last_line.startswith('-') or last_line.endswith(')')): + return True + if schema == 'new': + summary_body = text.split('## Summary', 1)[1].split('## More', 1)[0] + if _LEGACY_USE_WHEN_RE.search(summary_body): + return True + if not all(marker in summary_body for marker in _SCHEMA_MARKERS): + return True + return False + + +_api_semaphore = threading.Semaphore(API_CONCURRENCY) +_api_bucket_lock = threading.Lock() +_api_tokens = [float(API_CONCURRENCY)] +_api_last_refill = [time.monotonic()] + + +def _api_throttle(): + """Token-bucket rate limiter: API_CONCURRENCY requests per API_MIN_INTERVAL*API_CONCURRENCY window.""" + _api_semaphore.acquire() + try: + with _api_bucket_lock: + now = time.monotonic() + elapsed = now - _api_last_refill[0] + refill = elapsed / API_MIN_INTERVAL + _api_tokens[0] = min(float(API_CONCURRENCY), _api_tokens[0] + refill) + _api_last_refill[0] = now + if _api_tokens[0] >= 1.0: + _api_tokens[0] -= 1.0 + else: + wait = (1.0 - _api_tokens[0]) * API_MIN_INTERVAL + _api_tokens[0] = 0.0 + time.sleep(wait) + finally: + _api_semaphore.release() + + +def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[str]: + _api_throttle() + trajectory = {'messages': prompt['messages']} + sp = SamplingParams(temperature=0.2, max_tokens=8192) + try: + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + except Exception as exc: + logger.warning(f'[api_fallback] error: {exc}') + return None + content = (reply.get('content') or '').strip() + if not content: + return None + m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) + if m: + content = m.group(1).strip() + return content + + +# ============================================================================= +# Core compression logic +# ============================================================================= + +def _extract_query_cot(row: Dict[str, Any]): + messages = row.get('messages') or [] + query, cot = '', '' + for m in messages: + if not isinstance(m, dict): + continue + role = m.get('role') or '' + if role == 'user' and not query: + query = (m.get('content') or '').strip() + elif role == 'assistant': + cot = (m.get('reasoning_content') or '').strip() + break + return query, cot + + +def _compress_batch_phase1( + rows: List[Dict[str, Any]], + condenser_sampler, + compress_params: SamplingParams, + special_tokens: set, + source_type: str, +) -> Optional[Dict[str, Any]]: + """Phase 1 (GPU): build prompts → vLLM sample → validate. Returns state for phase 2.""" + _MAX_COT_CHARS = 30_000 + + if source_type == 'hard': + return _compress_hard_phase1(rows, condenser_sampler, compress_params, + special_tokens, source_type) + + prompts: List[Optional[Dict[str, Any]]] = [] + meta: List[Dict[str, Any]] = [] + for i, row in enumerate(rows): + query, cot = _extract_query_cot(row) + if not query or len(cot) < MIN_TEXT_CHARS or len(cot) > _MAX_COT_CHARS: + continue + schema = 'legacy' if (i % 2 == 0) else 'new' + q_hint = EMBED_QUERY_Q_LEGACY if schema == 'legacy' else EMBED_QUERY_Q + c_hint = EMBED_QUERY_COT_LEGACY if schema == 'legacy' else EMBED_QUERY_COT + + if len(query) < MIN_TEXT_CHARS: + prompts.append(None) + else: + user = COMPRESS_USER.format(query=q_hint, text=query) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user}, + ]}) + user_c = COMPRESS_USER.format(query=c_hint, text=cot) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_c}, + ]}) + meta.append({'query_raw': query, 'cot_raw': cot, 'schema': schema, + 'q_hint': q_hint, 'source': source_type, + 'row_id': row.get('id', str(i))}) + + if not prompts: + return {'final': []} + + sampler_input = [p for p in prompts if p is not None] + sampler_pos = [ri for ri, p in enumerate(prompts) if p is not None] + try: + sampler_responses = condenser_sampler.sample(sampler_input, compress_params) + except Exception as exc: + logger.warning(f'[compress] sampler error: {exc}') + sampler_responses = [None] * len(sampler_input) + + responses = [None] * len(prompts) + for resp, pos in zip(sampler_responses, sampler_pos): + responses[pos] = resp + + decoded: List[str] = [] + fallback_indices: List[int] = [] + for ri in range(len(prompts)): + pair_idx = ri // 2 + schema = meta[pair_idx]['schema'] + if prompts[ri] is None: + decoded.append(meta[pair_idx]['query_raw']) + continue + resp = responses[ri] + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + for tok in special_tokens: + text = text.replace(tok, '') + text = text.rstrip() + if not _is_truncated_compression(text, schema): + decoded.append(text) + else: + decoded.append('') + fallback_indices.append(ri) + + return {'prompts': prompts, 'meta': meta, 'decoded': decoded, + 'fallback_indices': fallback_indices} + + +def _compress_batch_phase2( + state: Dict[str, Any], + api_client: OpenAIClient, +) -> List[Dict[str, Any]]: + """Phase 2 (no GPU): API fallback → build results.""" + if 'final' in state: + return state['final'] + + prompts = state['prompts'] + decoded = state['decoded'] + fallback_indices = state['fallback_indices'] + is_hard = state.get('hard', False) + meta = state.get('meta') # None for hard + + # Track which prompts used API fallback + api_set: set = set() + if fallback_indices: + api_futures = {} + with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: + for ri in fallback_indices: + api_futures[pool.submit(_api_compress, api_client, prompts[ri])] = ri + for fut in as_completed(api_futures): + ri = api_futures[fut] + api_result = fut.result() + schema = 'new' if is_hard else meta[ri // 2]['schema'] + if api_result and not _is_truncated_compression(api_result, schema): + decoded[ri] = api_result + api_set.add(ri) + + state['api_set'] = api_set + if is_hard: + return _build_hard_results(state) + return _build_think_index_results(state) + + +def _build_think_index_results(state: Dict[str, Any]) -> List[Dict[str, Any]]: + meta = state['meta'] + decoded = state['decoded'] + api_set = state.get('api_set', set()) + results = [] + for pair_idx in range(len(meta)): + q_text = decoded[pair_idx * 2] + c_text = decoded[pair_idx * 2 + 1] + if not q_text or not c_text: + continue + q_method = 'api' if (pair_idx * 2) in api_set else 'vllm' + c_method = 'api' if (pair_idx * 2 + 1) in api_set else 'vllm' + results.append({ + 'anchor_text': q_text, + 'positive_text': c_text, + 'negative_texts': [], + 'source': meta[pair_idx]['source'], + 'query_raw': meta[pair_idx]['query_raw'], + 'cot_raw': meta[pair_idx]['cot_raw'], + 'anchor_method': q_method, + 'positive_method': c_method, + }) + return results + + +def _build_hard_results(state: Dict[str, Any]) -> List[Dict[str, Any]]: + group_sizes = state['group_sizes'] + decoded = state['decoded'] + source_type = state['source_type'] + raw_groups = state['raw_groups'] + api_set = state.get('api_set', set()) + results = [] + offset = 0 + for gi, gs in enumerate(group_sizes): + q_text = decoded[offset] + c_text = decoded[offset + 1] + if not q_text or not c_text: + offset += gs + continue + neg_texts = [] + neg_raws = [] + neg_methods = [] + for ni in range(2, gs): + nt = decoded[offset + ni] + if nt: + neg_texts.append(nt) + neg_raws.append(raw_groups[gi]['negs_raw'][ni - 2]) + neg_methods.append('api' if (offset + ni) in api_set else 'vllm') + q_method = 'api' if offset in api_set else 'vllm' + c_method = 'api' if (offset + 1) in api_set else 'vllm' + results.append({ + 'anchor_text': q_text, + 'positive_text': c_text, + 'negative_texts': neg_texts, + 'source': source_type, + 'query_raw': raw_groups[gi]['query_raw'], + 'cot_raw': raw_groups[gi]['cot_raw'], + 'negs_raw': neg_raws, + 'anchor_method': q_method, + 'positive_method': c_method, + 'neg_methods': neg_methods, + }) + offset += gs + return results + + +def _compress_hard_phase1( + rows: List[Dict[str, Any]], + condenser_sampler, + compress_params: SamplingParams, + special_tokens: set, + source_type: str, +) -> Dict[str, Any]: + """Phase 1 for hard rows: vLLM sample + validate. Returns state for phase 2.""" + _MAX_COT_CHARS = 30_000 + + prompts: List[Dict[str, Any]] = [] + group_sizes: List[int] = [] + row_ids: List[str] = [] + raw_groups: List[Dict[str, Any]] = [] + + for row in rows: + query, cot = _extract_query_cot(row) + if not query or not cot or len(cot) > _MAX_COT_CHARS: + continue + negatives = row.get('negatives') or [] + valid_negs = [n for n in negatives + if n and len(n) <= _MAX_COT_CHARS] + + user_q = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_Q, text=query) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_q}, + ]}) + user_c = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_COT, text=cot) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_c}, + ]}) + for neg in valid_negs: + user_n = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_COT, text=neg) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_n}, + ]}) + group_sizes.append(2 + len(valid_negs)) + row_ids.append(row.get('id', '')) + raw_groups.append({'query_raw': query, 'cot_raw': cot, 'negs_raw': valid_negs}) + + if not prompts: + return {'hard': True, 'prompts': [], 'group_sizes': [], 'row_ids': [], + 'decoded': [], 'fallback_indices': [], 'source_type': source_type, + 'raw_groups': []} + + try: + responses = condenser_sampler.sample(prompts, compress_params) + except Exception as exc: + logger.warning(f'[compress-hard] sampler error: {exc}') + responses = [None] * len(prompts) + + decoded: List[str] = [] + fallback_indices: List[int] = [] + for ri, resp in enumerate(responses): + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + for tok in special_tokens: + text = text.replace(tok, '') + text = text.rstrip() + if text and not _is_truncated_compression(text, 'new'): + decoded.append(text) + else: + decoded.append('') + fallback_indices.append(ri) + + return {'hard': True, 'prompts': prompts, 'group_sizes': group_sizes, + 'row_ids': row_ids, 'decoded': decoded, + 'fallback_indices': fallback_indices, 'source_type': source_type, + 'raw_groups': raw_groups} + + +# ============================================================================= +# Main pipeline +# ============================================================================= + +def main(): + device_groups = [ + DeviceGroup(name='condenser_sampler', + ranks=list(range(CONDENSER_GPUS)), + device_type='GPU'), + ] + condenser_mesh = DeviceMesh.from_sizes( + world_size=CONDENSER_GPUS, dp_size=CONDENSER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=CONDENSER_GPUS, groups=device_groups) + + # -- Load raw datasets ---------------------------------------------------- + from datasets import Dataset as HFDataset + + dataset_think = get_dataset_think(total=TOTAL_SAMPLES, load_from_cache_file=True) + if THINK_CAP and len(dataset_think.dataset) > THINK_CAP: + dataset_think.dataset = dataset_think.dataset.select(range(THINK_CAP)) + ds_think = dataset_think.dataset + logger.info(f'[load] think={len(ds_think)}') + + ds_index_obj = get_dataset_index(total=None, load_from_cache_file=True) + ds_index = ds_index_obj.dataset + if INDEX_CAP and len(ds_index) > INDEX_CAP: + ds_index = ds_index.select(range(INDEX_CAP)) + logger.info(f'[load] index={len(ds_index)}') + + ds_hard_raw = get_dataset_hard(max_negatives=HARD_MAX_NEGATIVES, load_from_cache_file=True) + if HARD_CAP and len(ds_hard_raw) > HARD_CAP: + ds_hard_raw = ds_hard_raw.select(range(HARD_CAP)) + n_hard = len(ds_hard_raw) + logger.info(f'[load] hard={n_hard}') + + # Convert hard to messages schema + hard_rows_list = [] + if n_hard > 0: + h_ids = ds_hard_raw['id'] + h_queries = ds_hard_raw['query'] + h_cots = ds_hard_raw['cot'] + h_responses = ds_hard_raw['response'] if 'response' in ds_hard_raw.column_names else [''] * n_hard + h_negatives = ds_hard_raw['negatives'] + for i in range(n_hard): + hard_rows_list.append({ + 'id': h_ids[i], + 'messages': [ + {'role': 'user', 'content': h_queries[i]}, + {'role': 'assistant', 'reasoning_content': h_cots[i], + 'content': h_responses[i] or ''}, + ], + 'negatives': h_negatives[i], + }) + + # Batch-convert HF Datasets to list-of-dicts + def _ds_to_rows(ds): + return [dict(zip(ds.column_names, vals)) for vals in zip(*(ds[c] for c in ds.column_names))] + + think_rows = _ds_to_rows(ds_think) + index_rows = _ds_to_rows(ds_index) + + # -- Setup condenser ------------------------------------------------------ + condenser_template = Qwen3_5Template( + model_id=CONDENSE_MODEL_ID, max_length=DATASET_MAX_TOKENS, + enable_thinking=False, truncation_strategy='delete') + special_tokens = set(condenser_template.tokenizer.all_special_tokens) + + condenser_sampler = vLLMSampler( + model_id=CONDENSE_MODEL_ID, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': COMPRESS_MAX_MODEL_LEN}, + device_mesh=condenser_mesh, + remote_group='condenser_sampler', + ) + condenser_sampler.set_template( + TEMPLATE_NAME, model_id=CONDENSE_MODEL_ID, enable_thinking=False, + truncation_strategy='delete', max_length=DATASET_MAX_TOKENS) + condenser_sampler._ray_get_timeout = SAMPLER_TIMEOUT + compress_params = SamplingParams( + max_tokens=8192, temperature=COMPRESS_TEMPERATURE, + top_p=COMPRESS_TOP_P, num_samples=1) + + api_client = OpenAIClient( + model=COMPRESS_MODEL, api_key=COMPRESS_API_KEY, base_url=COMPRESS_BASE_URL) + + # -- Resume support ---------------------------------------------------------- + os.makedirs(OUTPUT_DIR, exist_ok=True) + progress = {'think': 0, 'index': 0, 'hard': 0} + if os.path.exists(PROGRESS_FILE): + with open(PROGRESS_FILE, 'r') as f: + progress = json.load(f) + logger.info(f'[resume] loaded progress: {progress}') + + _results_lock = PosixFileLock(RESULTS_JSONL + '.lock') + + def _flush_results(records: List[Dict[str, Any]]): + if not records: + return + lines = [json.dumps(r, ensure_ascii=False) + '\n' for r in records] + with _results_lock: + with open(RESULTS_JSONL, 'a', encoding='utf-8') as f: + f.writelines(lines) + + def _save_progress(): + tmp = PROGRESS_FILE + '.tmp' + with open(tmp, 'w') as f: + json.dump(progress, f) + os.replace(tmp, PROGRESS_FILE) + + # -- Process in batches (pipelined: vLLM batch N+1 overlaps API fallback N) - + total_flushed = 0 + if os.path.exists(RESULTS_JSONL): + with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: + total_flushed = sum(1 for l in f if l.strip()) + if total_flushed: + logger.info(f'[resume] {total_flushed} records already in results.jsonl') + + def _process_source(rows, source_type, label): + nonlocal total_flushed + n_total = len(rows) + skip = progress.get(source_type, 0) + if skip >= n_total: + logger.info(f'[{label}] skipped (already done {skip}/{n_total})') + return + if skip > 0: + logger.info(f'[{label}] resuming from row {skip}/{n_total}') + + bg_pool = ThreadPoolExecutor(max_workers=1) + pending = None # (future, batch_start, batch_len) + + def _drain_pending(): + nonlocal total_flushed, pending + if pending is None: + return + fut, p_start, p_len = pending + batch_results = fut.result() + _flush_results(batch_results) + total_flushed += len(batch_results) + progress[source_type] = p_start + p_len + _save_progress() + pending = None + + for start in range(skip, n_total, BATCH_SIZE): + batch = rows[start:start + BATCH_SIZE] + state = _compress_batch_phase1( + batch, condenser_sampler, compress_params, + special_tokens, source_type) + _drain_pending() + pending = ( + bg_pool.submit(_compress_batch_phase2, state, api_client), + start, len(batch)) + n_done = start + len(batch) + if n_done % (BATCH_SIZE * 10) == 0 or n_done >= n_total: + logger.info(f'[{label}] {n_done}/{n_total} vLLM done, ' + f'{total_flushed} records flushed (last batch pending)') + + _drain_pending() + bg_pool.shutdown(wait=False) + logger.info(f'[{label}] complete, {total_flushed} total records flushed') + + _process_source(hard_rows_list, 'hard', 'hard') + _process_source(think_rows, 'think', 'think') + _process_source(index_rows, 'index', 'index') + + # -- Convert JSONL → HF Dataset ------------------------------------------- + logger.info(f'[save] converting results.jsonl to HF Dataset...') + all_results = [] + with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: + for line_no, line in enumerate(f, 1): + if not line.strip(): + continue + try: + all_results.append(json.loads(line)) + except json.JSONDecodeError: + logger.warning(f'[save] skipping malformed line {line_no} (truncated resume?)') + logger.info(f'[save] total records: {len(all_results)}') + out_ds = HFDataset.from_dict({ + 'anchor_text': [r['anchor_text'] for r in all_results], + 'positive_text': [r['positive_text'] for r in all_results], + 'negative_texts': [r['negative_texts'] for r in all_results], + 'source': [r['source'] for r in all_results], + 'query_raw': [r.get('query_raw', '') for r in all_results], + 'cot_raw': [r.get('cot_raw', '') for r in all_results], + 'negs_raw': [r.get('negs_raw', []) for r in all_results], + }) + out_ds.save_to_disk(OUTPUT_DIR + '/dataset') + logger.info(f'[save] dataset saved to {OUTPUT_DIR}/dataset') + logger.info(f'[stats] think={sum(1 for r in all_results if r["source"]=="think")} ' + f'index={sum(1 for r in all_results if r["source"]=="index")} ' + f'hard={sum(1 for r in all_results if r["source"]=="hard")}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py index bc69c56f..97ab3b12 100644 --- a/cookbook/exp/embedding/train_embedding_full_ddp.py +++ b/cookbook/exp/embedding/train_embedding_full_ddp.py @@ -1,277 +1,60 @@ -"""LoRA embedding training with online compression via frozen vLLM condenser. +"""Full-parameter embedding training on pre-compressed dataset. -Architecture (8 GPUs total): - - Ranks 0-3 (``model``): Trainable embedding model with LoRA, InfoNCE loss. - - Ranks 4-7 (``condenser_sampler``): Frozen vLLM condenser for online compression. +Reads the pre-compressed HF Dataset produced by make_embedding_dataset.py, +encodes features, trains with InfoNCE loss. -When the condenser sampler truncates or regresses to the legacy schema, an -external OpenAI-compatible API produces the correct compression. The failure is -logged to failures.jsonl for offline SFT data regeneration. +Architecture (4 GPUs): + - Ranks 0-3: Trainable embedding model, InfoNCE loss. Launch: - python cookbook/exp/train_embedding_lora_ddp.py + python cookbook/exp/embedding/train_embedding_full_ddp.py """ -import hashlib -import json import os -import re -import sys -import threading import time -from concurrent.futures import ThreadPoolExecutor -from pathlib import Path from typing import Any, Dict, List, Literal, Optional import swanlab import twinkle from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger -from twinkle.data_format import SamplingParams -from twinkle.dataloader import DataLoader from twinkle.loss import InfonceLoss from twinkle.metric import EmbeddingMetric from twinkle.model import TransformersModel from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler from twinkle.template import Qwen3_5Template, Template -from twinkle.utils.parallel import PosixFileLock -from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient - -sys.path.insert(0, str(Path(__file__).resolve().parent)) -from dataset_think import get_dataset as get_dataset_think # noqa: E402 -from dataset_index import get_dataset as get_dataset_index # noqa: E402 logger = get_logger() # -- Backend selection -------------------------------------------------------- BACKEND: Literal['transformers', 'megatron'] = 'transformers' -# Condenser (online compression + LoRA self-improvement); embedding model trains LoRA on top of MODEL_ID. -CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') -TEMPLATE_NAME = 'Qwen3_5Template' -# -- GPU placement (8 total) -------------------------------------------------- -MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) -CONDENSER_SAMPLER_GPUS = int(os.environ.get('CONDENSER_SAMPLER_GPUS', 4)) -NUM_GPUS = MODEL_GPUS + CONDENSER_SAMPLER_GPUS +# -- GPU placement ------------------------------------------------------------ +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 8)) # -- Embedding training hyper-params ------------------------------------------ EMB_MAX_LENGTH = 8192 HARD_NEGATIVES = None -# 0.07 keeps gradient on diag pairs until cosine clears ~0.75; 0.03 saturated near 0.40. TEMPERATURE = 0.07 -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 32)) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 64)) LEARNING_RATE = 1e-5 GRADIENT_ACCUMULATION_STEPS = 1 LOG_INTERVAL = 2 SAVE_INTERVAL = 2000 NUM_EPOCHS = 1 -TOTAL_SAMPLES: Optional[int] = None -# Post-build caps on each loader (None = no cap). Applied via .select() before mix. -THINK_CAP: Optional[int] = 400_000 -INDEX_CAP: Optional[int] = 400_000 +# -- Dataset path (output of make_embedding_dataset.py) ----------------------- +DATASET_PATH = os.environ.get('EMB_DATASET_PATH', 'ms://twinkle-kit/qth-embedding') MIX_SHUFFLE_SEED = 42 # -- Resume from checkpoint --------------------------------------------------- -# Empty by default — build_model falls back to MODEL_ID (the published emb model). -# Set both to point at a local in-progress run only when resuming the *same* schedule. RESUME_CHECKPOINT = os.environ.get('RESUME_CHECKPOINT', '') RESUME_STEP = int(os.environ.get('RESUME_STEP', 0)) -# -- Online-compression knobs ------------------------------------------------- -# Below this length, condenser fabricates content for open-ended short prompts; -# query passes through as qr verbatim and cot rows are dropped from training. -MIN_TEXT_CHARS = 256 -DATASET_MAX_TOKENS = 32768 -COMPRESS_TEMPERATURE = 0.2 -COMPRESS_TOP_P = 0.5 -COMPRESS_MAX_MODEL_LEN = 32768 - -# How many BATCH_SIZE chunks to fetch and compress in one vLLM call. -PREFETCH_BATCH_MULTIPLIER = int(os.environ.get('PREFETCH_BATCH_MULTIPLIER', 8)) - -# -- OpenAI API fallback for truncated compressions --------------------------- -COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '') -COMPRESS_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') -COMPRESS_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') -# Minimum gap between API calls (seconds); bounds dashscope qps under provider limits. -API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) -API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 8)) -# vLLM sampler timeout (seconds); if a sample() call exceeds this, fall back to API. -SAMPLER_TIMEOUT = float(os.environ.get('SAMPLER_TIMEOUT', 300)) - -# -- Output paths ------------------------------------------------------------- -OUTPUT_DIR = f'./output/embedding_lora_{BACKEND}' -RESPONSE_LOG = os.environ.get('RESPONSE_LOG', f'./output/embedding_lora_{BACKEND}/responses.jsonl') -FAILURE_LOG = os.environ.get('FAILURE_LOG', f'./output/embedding_lora_{BACKEND}/failures.jsonl') - - -# ============================================================================= -# Prompts (from make_condenser_dataset.py — "## Summary" format) -# ============================================================================= - -COMPRESS_SYSTEM = """\ -You are a compression and summary assistant. For the (query, source) pair, emit a Markdown \ -answer with TWO sections, designed to pair with the `extract_compressed` tool: \ -the reader absorbs `## Summary` directly, then calls `extract_compressed` \ -on any topic-key listed under `## More` to recover its \ -fuller content. - - `## Summary` — extreme-density text the reader reads directly. - `## More` — a topic index whose keys are valid arguments \ -to `extract_compressed` for recovering material not captured inline. - -Together the two sections must form a COMPLETE, NON-DISTORTING inventory of the \ -source for the query — nothing essential lost, nothing implied that the source \ -does not support. NO preamble, NO meta-commentary, NO code fences wrapping the \ -whole output. - -Output skeleton: - -## Summary -Topic: - - -## More -- : -- ... - -Format selection for the inline body (pick the MOST COMPACT form per query, mix \ -when helpful): -- Interface / signature → code notation directly: `func(a:int)->str` -- Factual / entity → telegraphic prose; drop function words; ":" for "is", "," \ -for "has" -- Skill / how-to / usage → lead with `Use when: `; numbered telegraphic \ -steps `1.do X 2.then Y`; close with `Output: ` when relevant -- Procedural → numbered short steps -- Analytical / design → hierarchical bullets with abbreviations - -`## Summary` rules: -1. TOPIC LINE — line 1 is ALWAYS `Topic: `, even when the \ -query is narrow. Anchors both the reader and the tool. -2. DENSITY — every token in the body carries query-relevant signal; cut filler. -3. PRIMARY-COMPLETE — never silently drop a fact essential to answering the \ -query. Anything cut for length MUST appear as a key under \ -`## More`. -4. NON-MISLEADING — phrasing must not let the reader infer anything the source \ -does not support; partial truths that mislead are worse than honest omissions \ -flagged in the index. -5. SELF-CONTAINED — the reader can act on the answer without re-opening the source. -6. FAITHFUL — only content the source supports; no fabrication, no extrapolation. -7. LANGUAGE — match the source language. -8. NO outer code fences around the whole answer; no meta-commentary. - -`## More` rules (MANDATORY — this section is never omitted): -1. FORMAT — each bullet is `- : `: - • topic-key — short, unambiguous, grounded in source vocabulary so the \ -`extract_compressed` tool can locate the aspect (e.g. `decorators`, \ -`error handling`, `pitfalls`). - • hint — tells WHAT the reader gains by expanding (concrete numbers, code \ -listings, secondary cases, edge details, related context, …); do NOT restate \ -the inline answer. -2. CRITERION — each bullet names an aspect that EXISTS in the source but is \ -NOT fully captured inline. Material that genuinely fits inline without \ -distortion MUST NOT be duplicated here. -3. FAITHFUL — hints must be grounded in the source; never speculate or invent. -4. ORDER — by relevance to the query, then by importance. -5. EMPTY CASE — if the source is so short / single-purpose that everything \ -fits inline, write a single line `- (none)`. - -Now begin.\ -""" - -COMPRESS_USER = ( - 'Downstream model will read your compressed block to decide whether to ' - 'expand it. Compress faithfully: preserve the passage topic + core facts. ' - 'Do NOT invent facts. Do NOT drop major facts. Do NOT write meta-commentary ' - 'about the Query (never write "Query info: absent", "no X mention", etc.); ' - 'if the passage does not address the Query, still summarize the passage. ' - 'CRITICAL LANGUAGE RULE: detect the dominant language of the Passage ' - '(NOT the Query, NOT this instruction) and write the ENTIRE output in that ' - 'same language; English passage → English output, Chinese passage → ' - 'Chinese output, Japanese passage → Japanese output. NEVER translate, ' - 'NEVER mix languages, NEVER copy these instructions into the output.\n\n' - '## Query (ordering hint only — still summarize the whole passage)\n{query}\n\n' - '## Passage\n{text}') - - -# ============================================================================= -# Logging helpers -# ============================================================================= - -_response_lock: Optional[PosixFileLock] = None -_failure_lock: Optional[PosixFileLock] = None - -# Monotonic global sample id; per-batch index would alias across batches. -_sample_counter = 0 -_sample_counter_lock = threading.Lock() - -_api_throttle_lock = threading.Lock() -_api_last_call = [0.0] - - -def _api_throttle(): - with _api_throttle_lock: - gap = time.monotonic() - _api_last_call[0] - if gap < API_MIN_INTERVAL: - time.sleep(API_MIN_INTERVAL - gap) - _api_last_call[0] = time.monotonic() - - -def _next_sample_id() -> int: - global _sample_counter - with _sample_counter_lock: - sid = _sample_counter - _sample_counter += 1 - return sid - - -def _log_responses(query_resp_text: str, cot_resp_text: str, idx: int, - query_raw: str = '', cot_raw: str = ''): - global _response_lock - if _response_lock is None: - os.makedirs(os.path.dirname(RESPONSE_LOG) or '.', exist_ok=True) - _response_lock = PosixFileLock(RESPONSE_LOG + '.lock') - - record = { - 'idx': idx, - 'query_raw': query_raw, - 'cot_raw': cot_raw, - 'query_compressed': query_resp_text, - 'cot_compressed': cot_resp_text, - } - line = json.dumps(record, ensure_ascii=False, default=str) + '\n' - with _response_lock: - with open(RESPONSE_LOG, 'a', encoding='utf-8') as f: - f.write(line) - - -def _log_failure(source_text: str, query: str, compressed: str, batch_idx: int): - global _failure_lock - if _failure_lock is None: - os.makedirs(os.path.dirname(FAILURE_LOG) or '.', exist_ok=True) - _failure_lock = PosixFileLock(FAILURE_LOG + '.lock') - - qhash = hashlib.md5(query.strip().encode('utf-8')).hexdigest()[:8] - record = { - 'id': f'{batch_idx}__{qhash}', - 'source': 'online_failure', - 'query': query, - 'original_len': len(source_text), - 'compressed_len': len(compressed), - 'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': COMPRESS_USER.format(query=query, text=source_text)}, - {'role': 'assistant', 'content': compressed}, - ], - } - line = json.dumps(record, ensure_ascii=False, default=str) + '\n' - with _failure_lock: - with open(FAILURE_LOG, 'a', encoding='utf-8') as f: - f.write(line) +# -- Output ------------------------------------------------------------------- +OUTPUT_DIR = f'./output/embedding_full_{BACKEND}' # ============================================================================= @@ -327,121 +110,9 @@ def save_checkpoint(model, name: str): # ============================================================================= -# Compression prompt building +# Feature encoding # ============================================================================= -# Hard-templated hints: the condenser SFT prior maps `Skill` to the legacy -# `Use when: / numbered steps / Output:` skeleton on long inputs; embedding the -# exact 4-line body template + explicit negative constraints is the only way to -# override it deterministically across query and cot sides. -EMBED_QUERY_Q = ( - 'Summarize this query for retrieval. ' - 'The body of ## Summary MUST follow this EXACT 4-line template — ' - 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' - 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the specific pattern, never generic labels.') -EMBED_QUERY_COT = ( - 'Summarize this reasoning trace for retrieval. ' - 'The body of ## Summary MUST follow this EXACT 4-line template — ' - 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' - 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the specific pattern, never generic labels.') - -# Legacy schema (Use when: / numbered steps / Output:) — mixed in 50/50 with the -# new schema to expose the embedder to schema-invariant semantic alignment. -# Both query and cot of the SAME pair always use the SAME schema; cross-schema -# anchors and positives would re-introduce the schema asymmetry we just fixed. -EMBED_QUERY_Q_LEGACY = ( - 'What problem does this passage address, and what skill or method is needed? ' - 'Topic must name the specific pattern, never generic labels. ' - 'Compress into a retrieval-friendly need description.') -EMBED_QUERY_COT_LEGACY = ( - 'Extract the reusable skill: trigger conditions, key steps, and expected output. ' - 'Topic names the method/pattern; format as "Use when: ...", numbered steps, ' - '"Output: ...". Compress into a standardized procedure for retrieval.') - - -def _extract_query_cot(row: Dict[str, Any]): - messages = row.get('messages') or [] - query, cot = '', '' - for m in messages: - if not isinstance(m, dict): - continue - role = m.get('role') or '' - if role == 'user' and not query: - query = (m.get('content') or '').strip() - elif role == 'assistant': - cot = (m.get('reasoning_content') or '').strip() - break - return query, cot - - -def _build_compress_prompts(rows: List[Dict[str, Any]]) -> tuple: - """Build prompts for compressing both query and cot per row. - - Returns (prompts, valid_indices, raw_pairs, prompt_queries, passthrough, schemas) - where: - - prompts: flat-interleaved [query_0, cot_0, query_1, cot_1, ...]; ``None`` means - passthrough (use raw text directly, do not call sampler) - - valid_indices: which rows passed the min-length filter - - raw_pairs: [(query, cot), ...] - - prompt_queries: the query string used for each prompt (for failure logging) - - passthrough: parallel to prompts; non-None text means "use this verbatim as qc" - - schemas: parallel to prompts; 'new' or 'legacy', drives validator branch - """ - prompts: List[Optional[Dict[str, Any]]] = [] - valid_indices: List[int] = [] - raw_pairs: List[tuple] = [] - prompt_queries: List[str] = [] - passthrough: List[Optional[str]] = [] - schemas: List[str] = [] - # Conservative char budget: 32768 max_length - 8192 gen - ~2k prompt overhead = ~22k tokens. - # 30k cap bounds vLLM batch latency (vLLM batches by max prompt length). - _MAX_COT_CHARS = 30_000 - for i, row in enumerate(rows): - query, cot = _extract_query_cot(row) - if not query or len(cot) < MIN_TEXT_CHARS: - continue - if len(cot) > _MAX_COT_CHARS: - continue - valid_indices.append(i) - raw_pairs.append((query, cot)) - # 50/50 schema mix; same schema for query+cot of one pair to keep alignment. - schema = 'legacy' if (i % 2 == 0) else 'new' - q_hint = EMBED_QUERY_Q_LEGACY if schema == 'legacy' else EMBED_QUERY_Q - c_hint = EMBED_QUERY_COT_LEGACY if schema == 'legacy' else EMBED_QUERY_COT - # Short query bypasses condenser to avoid skeleton-induced hallucination. - if len(query) < MIN_TEXT_CHARS: - prompts.append(None) - passthrough.append(query) - else: - user = COMPRESS_USER.format(query=q_hint, text=query) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user}, - ]}) - passthrough.append(None) - prompt_queries.append(q_hint) - schemas.append(schema) - user = COMPRESS_USER.format(query=c_hint, text=cot) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user}, - ]}) - prompt_queries.append(c_hint) - passthrough.append(None) - schemas.append(schema) - return prompts, valid_indices, raw_pairs, prompt_queries, passthrough, schemas - - def _get_first_feature(decoded_text: str, template: Template, role: str) -> Optional[Dict[str, Any]]: if not decoded_text: return None @@ -450,77 +121,42 @@ def _get_first_feature(decoded_text: str, template: Template, role: str) -> Opti {'role': 'user', 'content': decoded_text}, {'role': 'assistant', 'content': 'Match the correct response here.'}, ]}) + if feat is None: + return None feat['labels'] = [1] else: feat = template.encode({'messages': [ {'role': 'user', 'content': 'Match the correct query here.'}, {'role': 'assistant', 'content': decoded_text}, ]}) + if feat is None: + return None feat['labels'] = [0] return feat -# ============================================================================= -# OpenAI API fallback -# ============================================================================= - -_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') -_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') - - -def _is_truncated_compression(text: str, schema: str = 'new') -> bool: - """Reject structurally incomplete OR schema-regressed condenser output. - - Triggers API fallback when the vLLM output: - * lacks ``## Summary`` / ``## More``, - * has an empty or unterminated ``## More`` bullet list, or - * (schema='new' only) regresses to the legacy ``Use when: / numbered-steps / - Output:`` skeleton instead of the mandated Problem/Skill/Knowledge 4-line - body — the dominant cot-side failure mode that drives sim < 0.45 drops on - the RAG index. - - For schema='legacy', body markers are intentionally NOT enforced: the legacy - template legitimately emits ``Use when:`` and the SFT prior already produces - that shape natively, so only structural completeness is checked. - """ - if not text or not text.strip(): - return True - if '## More' not in text or '## Summary' not in text: - return True - after_more = text.split('## More', 1)[1].strip() - if not after_more: - return True - last_line = after_more.splitlines()[-1].strip() - if not (last_line.startswith('-') or last_line.endswith(')')): - return True - if schema == 'new': - summary_body = text.split('## Summary', 1)[1].split('## More', 1)[0] - if _LEGACY_USE_WHEN_RE.search(summary_body): - return True - if not all(marker in summary_body for marker in _SCHEMA_MARKERS): - return True - return False - - -def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[str]: - """Call external API to compress when vLLM truncates.""" - _api_throttle() - trajectory = {'messages': prompt['messages']} - # Cap max_tokens to leave ample prompt headroom inside the API model context. - sp = SamplingParams(temperature=0.2, max_tokens=8192) - try: - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - except Exception as exc: - logger.warning(f'[api_fallback] error: {exc}') - return None - content = (reply.get('content') or '').strip() - if not content: - return None - # Strip outer code fence if present - m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) - if m: - content = m.group(1).strip() - return content +def _encode_batch( + rows: List[Dict[str, Any]], + emb_template: Template, +) -> List[Dict[str, Any]]: + """Encode pre-compressed texts into embedding features.""" + features: List[Dict[str, Any]] = [] + for row in rows: + anchor_text = row['anchor_text'] + positive_text = row['positive_text'] + negative_texts = row.get('negative_texts') or [] + + feat_q = _get_first_feature(anchor_text, emb_template, role='anchor') + feat_c = _get_first_feature(positive_text, emb_template, role='positive') + if not feat_q or not feat_c: + continue + features.append(feat_q) + features.append(feat_c) + for neg_text in negative_texts: + feat_neg = _get_first_feature(neg_text, emb_template, role='positive') + if feat_neg: + features.append(feat_neg) + return features # ============================================================================= @@ -528,42 +164,27 @@ def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[ # ============================================================================= def train(): - # -------- Device groups (2 groups) ---------------------------------------- device_groups = [ DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), - DeviceGroup(name='condenser_sampler', - ranks=list(range(MODEL_GPUS, MODEL_GPUS + CONDENSER_SAMPLER_GPUS)), - device_type='GPU'), ] model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) - condenser_sampler_mesh = DeviceMesh.from_sizes( - world_size=CONDENSER_SAMPLER_GPUS, dp_size=CONDENSER_SAMPLER_GPUS) - - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups) - - # -------- Data ----------------------------------------------------------- - dataset = get_dataset_think(total=TOTAL_SAMPLES, load_from_cache_file=True) - if THINK_CAP and len(dataset.dataset) > THINK_CAP: - dataset.dataset = dataset.dataset.select(range(THINK_CAP)) - if INDEX_CAP != 0: - from datasets import concatenate_datasets - ds_index = get_dataset_index(total=None, load_from_cache_file=True) - if INDEX_CAP and len(ds_index.dataset) > INDEX_CAP: - ds_index.dataset = ds_index.dataset.select(range(INDEX_CAP)) - n_think = len(dataset.dataset) - n_index = len(ds_index.dataset) - # Both loaders emit identical {id, source, messages} schema post-QP. - dataset.dataset = concatenate_datasets( - [dataset.dataset, ds_index.dataset]).shuffle(seed=MIX_SHUFFLE_SEED) - logger.info(f'[mix] think={n_think} + index={n_index} → total={len(dataset.dataset)}') - _mega_batch_size = BATCH_SIZE * PREFETCH_BATCH_MULTIPLIER - dataloader = DataLoader(dataset=dataset, batch_size=_mega_batch_size, shuffle=True) - total_forward_steps = len(dataloader) * PREFETCH_BATCH_MULTIPLIER * NUM_EPOCHS - optimizer_steps = total_forward_steps // GRADIENT_ACCUMULATION_STEPS - - # -------- Embedding model (4 GPU) ---------------------------------------- + twinkle.initialize(mode='ray', nproc_per_node=MODEL_GPUS, groups=device_groups) + + # -- Load pre-compressed dataset ------------------------------------------ + from twinkle.dataset import Dataset as TwinkleDataset, DatasetMeta + logger.info(f'[data] loading pre-compressed dataset from {DATASET_PATH}') + dataset = TwinkleDataset(DatasetMeta(dataset_id=DATASET_PATH), download_mode='force_redownload') + dataset = dataset.dataset.shuffle(seed=MIX_SHUFFLE_SEED) + logger.info(f'[data] {len(dataset)} rows loaded') + + # -- Compute steps -------------------------------------------------------- + rows_per_step = BATCH_SIZE + total_steps = (len(dataset) // rows_per_step) * NUM_EPOCHS + optimizer_steps = total_steps // GRADIENT_ACCUMULATION_STEPS + + # -- Model ---------------------------------------------------------------- model = build_model(model_mesh) model.set_processor(InputProcessor) model.set_loss(InfonceLoss, temperature=TEMPERATURE, use_batch=True, @@ -571,264 +192,78 @@ def train(): setup_optimizer(model, optimizer_steps) model.add_metric(EmbeddingMetric, is_training=True) - # -------- Condenser sampler (4 GPU, vLLM) -------------------------------- - emb_template = Qwen3_5Template(model_id=MODEL_ID, max_length=EMB_MAX_LENGTH, enable_thinking=False) - # Special tokens come from the condenser tokenizer because the leak we strip is in its decoded output. - condenser_template = Qwen3_5Template(model_id=CONDENSE_MODEL_ID, max_length=DATASET_MAX_TOKENS, - enable_thinking=False) - _special_tokens = set(condenser_template.tokenizer.all_special_tokens) - condenser_sampler = vLLMSampler( - model_id=CONDENSE_MODEL_ID, - engine_args={ - 'gpu_memory_utilization': 0.8, - 'max_model_len': COMPRESS_MAX_MODEL_LEN, - }, - device_mesh=condenser_sampler_mesh, - remote_group='condenser_sampler', - ) - condenser_sampler.set_template( - TEMPLATE_NAME, model_id=CONDENSE_MODEL_ID, enable_thinking=False, - truncation_strategy='delete', max_length=DATASET_MAX_TOKENS) - compress_params = SamplingParams( - max_tokens=8192, - temperature=COMPRESS_TEMPERATURE, - top_p=COMPRESS_TOP_P, - num_samples=1, - ) - - condenser_sampler._ray_get_timeout = SAMPLER_TIMEOUT - _sampler_epoch = 0 - - def _rebuild_sampler(): - """Kill stuck actors and recreate the vLLM sampler from scratch.""" - nonlocal condenser_sampler, _sampler_epoch - import ray - for actor in getattr(condenser_sampler, '_actors', []): - try: - ray.kill(actor, no_restart=True) - except Exception: - pass - logger.warning('[sampler] killed stuck actors, recreating sampler \u2026') - new = vLLMSampler( - model_id=CONDENSE_MODEL_ID, - engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': COMPRESS_MAX_MODEL_LEN}, - device_mesh=condenser_sampler_mesh, - remote_group='condenser_sampler', - ) - new.set_template( - TEMPLATE_NAME, model_id=CONDENSE_MODEL_ID, enable_thinking=False, - truncation_strategy='delete', max_length=DATASET_MAX_TOKENS) - new._ray_get_timeout = SAMPLER_TIMEOUT - condenser_sampler = new - _sampler_epoch += 1 - logger.warning('[sampler] sampler rebuilt successfully') - - # -------- OpenAI API client for fallback --------------------------------- - api_client = OpenAIClient( - model=COMPRESS_MODEL, - api_key=COMPRESS_API_KEY, - base_url=COMPRESS_BASE_URL, - ) + emb_template = Qwen3_5Template( + model_id=MODEL_ID, max_length=EMB_MAX_LENGTH, + enable_thinking=False, truncation_strategy='delete') logger.info(get_device_placement()) logger.info(model.get_train_configs()) - logger.info(f'Total forward steps: {total_forward_steps}, optimizer steps: {optimizer_steps}') - if RESUME_STEP > 0: - logger.info(f'Resuming from step {RESUME_STEP}, checkpoint: {RESUME_CHECKPOINT}') - logger.info(f'Starting at epoch {RESUME_STEP // (total_forward_steps // NUM_EPOCHS)}, ' - f'skipping {RESUME_STEP - (RESUME_STEP // (total_forward_steps // NUM_EPOCHS)) * (total_forward_steps // NUM_EPOCHS)} batches') + logger.info(f'Total steps: {total_steps}, optimizer steps: {optimizer_steps}') swanlab.init(project='twinkle', config={ 'backend': BACKEND, 'model_id': MODEL_ID, - 'condense_model_id': CONDENSE_MODEL_ID, 'batch_size': BATCH_SIZE, 'lr': LEARNING_RATE, 'temperature': TEMPERATURE, 'emb_max_length': EMB_MAX_LENGTH, - 'DATASET_MAX_TOKENS': DATASET_MAX_TOKENS, + 'dataset_path': DATASET_PATH, }) - # -------- Train loop ----------------------------------------------------- - def _sample_batch(raw_batch): - """Compress via vLLM sampler; fall back to API on truncation.""" - _t_enter = time.monotonic() - compress_prompts, valid_indices, raw_pairs, prompt_queries, passthrough, schemas = \ - _build_compress_prompts(raw_batch) - _t_build = time.monotonic() - if len(compress_prompts) < 4: - return None + # -- Train loop ----------------------------------------------------------- + cur_step = 0 + _skip_rows = RESUME_STEP * rows_per_step # approximate rows to skip - # Only submit non-passthrough prompts to the sampler. - sampler_input = [p for p in compress_prompts if p is not None] - sampler_pos = [ri for ri, p in enumerate(compress_prompts) if p is not None] - if sampler_input: - try: - sampler_responses = condenser_sampler.sample(sampler_input, compress_params) - except Exception as exc: - logger.warning(f'[sampler] error \u2192 API fallback: {exc}') - sampler_responses = [None] * len(sampler_input) - if 'Timeout' in type(exc).__name__: - try: - _rebuild_sampler() - except Exception as re_exc: - logger.error(f'[sampler] rebuild failed: {re_exc}') - else: - sampler_responses = [] - _t_sample = time.monotonic() - - responses = [None] * len(compress_prompts) - for resp, pos in zip(sampler_responses, sampler_pos): - responses[pos] = resp - - # Extract decoded texts; detect truncations and fall back to API - decoded_texts: List[Optional[str]] = [None] * len(compress_prompts) - fallback_indices: List[int] = [] - for ri in range(len(compress_prompts)): - if passthrough[ri] is not None: - decoded_texts[ri] = passthrough[ri] + for epoch in range(NUM_EPOCHS): + for start in range(0, len(dataset), rows_per_step): + if start < _skip_rows: continue - resp = responses[ri] - seq = resp.sequences[0] if resp and resp.sequences else None - text = '' - if seq and seq.stop_reason != 'length' and seq.decoded: - text = seq.decoded - for tok in _special_tokens: - text = text.replace(tok, '') - text = text.rstrip() - - needs_fallback = (not seq or seq.stop_reason == 'length' - or _is_truncated_compression(text, schemas[ri])) - if not needs_fallback: - decoded_texts[ri] = text - else: - fallback_indices.append(ri) - - _api_calls = len(fallback_indices) - if fallback_indices: - from concurrent.futures import as_completed - api_futures = {} - with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as api_pool: - for ri in fallback_indices: - api_futures[api_pool.submit(_api_compress, api_client, compress_prompts[ri])] = ri - for fut in as_completed(api_futures): - ri = api_futures[fut] - api_result = fut.result() - if api_result and not _is_truncated_compression(api_result, schemas[ri]): - decoded_texts[ri] = api_result - pair_idx = ri // 2 - q_raw, c_raw = raw_pairs[pair_idx] - source_text = q_raw if ri % 2 == 0 else c_raw - _log_failure(source_text, prompt_queries[ri], api_result, - valid_indices[pair_idx]) - else: - decoded_texts[ri] = '' - _t_api = time.monotonic() - - # Build embedding features from decoded texts - emb_features: List[Dict[str, Any]] = [] - for i in range(0, len(decoded_texts), 2): - q_text = decoded_texts[i] - c_text = decoded_texts[i + 1] - q_raw, c_raw = raw_pairs[i // 2] - _log_responses(q_text, c_text, _next_sample_id(), - query_raw=q_raw, cot_raw=c_raw) - feat_q = _get_first_feature(q_text, emb_template, role='anchor') - feat_c = _get_first_feature(c_text, emb_template, role='positive') - if feat_q and feat_c: - emb_features.append(feat_q) - emb_features.append(feat_c) - _t_feat = time.monotonic() - - logger.info( - f'[prefetch] prompts={len(sampler_input)} api={_api_calls} feats={len(emb_features)} | ' - f'build={_t_build - _t_enter:.1f}s ' - f'vllm={_t_sample - _t_build:.1f}s ' - f'api={_t_api - _t_sample:.1f}s feat={_t_feat - _t_api:.1f}s ' - f'total={_t_feat - _t_enter:.1f}s') - - _target = BATCH_SIZE * 2 - minibatches = [emb_features[i:i + _target] for i in range(0, len(emb_features), _target)] - minibatches = [mb for mb in minibatches if len(mb) >= 4] - return minibatches if minibatches else None - - cur_step = RESUME_STEP - _batches_per_epoch = len(dataloader) - _steps_per_mega = PREFETCH_BATCH_MULTIPLIER - _start_epoch = cur_step // (_batches_per_epoch * _steps_per_mega) if cur_step > 0 else 0 - _skip_batches_in_epoch = max(0, cur_step // _steps_per_mega - _start_epoch * _batches_per_epoch) - - _ema_prefetch = 0.0 - _ema_train = 0.0 - _ema_alpha = 0.1 - - prefetch_executor = ThreadPoolExecutor(max_workers=1) - for epoch in range(_start_epoch, NUM_EPOCHS): - if _skip_batches_in_epoch > 0: - dataloader.skip_consumed_samples(_skip_batches_in_epoch * _mega_batch_size) - batch_iter = iter(dataloader) - _skip_batches_in_epoch = 0 - - first = next(batch_iter, None) - future = prefetch_executor.submit(_sample_batch, first) if first else None - - for raw_mega_batch in batch_iter: + + batch_rows = dataset[start:start + rows_per_step] + # HF Dataset slicing returns dict of lists; convert to list of dicts + n_rows = len(batch_rows['anchor_text']) + rows_list = [{k: batch_rows[k][i] for k in batch_rows} + for i in range(n_rows)] + t0 = time.monotonic() - minibatches = future.result() if future else None - t_prefetch = time.monotonic() - t0 - future = prefetch_executor.submit(_sample_batch, raw_mega_batch) + features = _encode_batch(rows_list, emb_template) + t_encode = time.monotonic() - t0 - if not minibatches: + if len(features) < 4: continue - for mb in minibatches: - t1 = time.monotonic() - model.forward_backward(inputs=mb, task='embedding') - model.clip_grad_and_step(gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - t_train = time.monotonic() - t1 - cur_step += 1 - - _ema_prefetch = _ema_alpha * t_prefetch + (1 - _ema_alpha) * _ema_prefetch if cur_step > RESUME_STEP + 1 else t_prefetch - _ema_train = _ema_alpha * t_train + (1 - _ema_alpha) * _ema_train if cur_step > RESUME_STEP + 1 else t_train - - if cur_step % LOG_INTERVAL == 0: - metric = model.calculate_metric(is_training=True) - _bottleneck = 'PREFETCH' if _ema_prefetch > _ema_train else 'TRAIN' - logger.info( - f'Epoch {epoch} Step {cur_step}/{total_forward_steps}, metric: {metric} | ' - f'prefetch={t_prefetch:.1f}s(ema {_ema_prefetch:.1f}) ' - f'train={t_train:.1f}s(ema {_ema_train:.1f}) ' - f'bottleneck={_bottleneck}') - log_dict = {} - for k, v in metric.items(): - if not v: - continue - try: - log_dict[k] = float(v) - except (ValueError, TypeError): - pass - log_dict['epoch'] = epoch - log_dict['prefetch_sec'] = round(t_prefetch, 2) - log_dict['train_sec'] = round(t_train, 2) - swanlab.log(log_dict, step=cur_step) - if cur_step % SAVE_INTERVAL == 0: - save_checkpoint(model, f'step_{cur_step}') - t_prefetch = 0.0 - - # Drain final mega-batch - if future: - minibatches = future.result() - future = None - if minibatches: - for mb in minibatches: - model.forward_backward(inputs=mb, task='embedding') - model.clip_grad_and_step(gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - cur_step += 1 - if cur_step % SAVE_INTERVAL == 0: - save_checkpoint(model, f'step_{cur_step}') - - prefetch_executor.shutdown(wait=False) + t1 = time.monotonic() + model.forward_backward(inputs=features, task='embedding') + model.clip_grad_and_step( + gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + t_train = time.monotonic() - t1 + cur_step += 1 + + if cur_step % LOG_INTERVAL == 0: + metric = model.calculate_metric(is_training=True) + logger.info( + f'Epoch {epoch} Step {cur_step}/{total_steps}, ' + f'metric: {metric} | ' + f'encode={t_encode:.2f}s train={t_train:.2f}s') + log_dict = {} + for k, v in metric.items(): + if not v: + continue + try: + log_dict[k] = float(v) + except (ValueError, TypeError): + pass + log_dict['epoch'] = epoch + log_dict['encode_sec'] = round(t_encode, 3) + log_dict['train_sec'] = round(t_train, 3) + swanlab.log(log_dict, step=cur_step) + if cur_step % SAVE_INTERVAL == 0: + save_checkpoint(model, f'step_{cur_step}') + save_checkpoint(model, 'last-checkpoint') + # Force sync: resolve any pending lazy remote calls (save) before exit + model.calculate_metric(is_training=True) + logger.info(f'Training complete. Final step: {cur_step}') if __name__ == '__main__': diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/rl/grpo.py new file mode 100644 index 00000000..7f43d93f --- /dev/null +++ b/cookbook/exp/rl/grpo.py @@ -0,0 +1,787 @@ +"""Pure GRPO training on AoPS dataset (no RAG, ablation baseline). + +Architecture (8 GPUs): + - 4 GPUs: sampler/rollout (vLLM TP=4) + - 4 GPUs: training model (FSDP) + +Pipeline per step: + 1. DataLoader yields a batch of math problems + 2. Sampler generates rollouts + 3. Reward (accuracy + format + gibberish) → GRPO advantage → model update + +Launch: + python cookbook/exp/rl/grpo.py +""" +import json +import os +import re +import random +from typing import Any, Dict, List, Tuple + +import numpy as np +import torch + +import twinkle +from twinkle import DeviceMesh, DeviceGroup, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.reward.base import Reward +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template + +logger = get_logger() + +# ============================================================================ +# Configuration +# ============================================================================ +MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') + +# GPU layout: 4 rollout + 4 train = 8 +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 4)) +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) +NUM_GPUS = SAMPLER_GPUS + MODEL_GPUS + +# Training hyperparams +NUM_GENERATIONS = int(os.environ.get('NUM_GENERATIONS', 8)) +MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 32768)) +LEARNING_RATE = float(os.environ.get('LR', 1e-5)) +MAX_STEPS = int(os.environ.get('MAX_STEPS', 5000)) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 8)) +MINI_BATCH_SIZE = int(os.environ.get('MINI_BATCH_SIZE', 8)) +MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 8)) +GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 1)) +SAVE_STEPS = int(os.environ.get('SAVE_STEPS', 100)) +ADV_CLIP = float(os.environ.get('ADV_CLIP', 1.0)) +LOSS_SPIKE_THRESHOLD = float(os.environ.get('LOSS_SPIKE_THRESHOLD', 10.0)) + +# Dataset +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +AOPS_SEED = int(os.environ.get('AOPS_SEED', 100)) + +# Output / diagnostics +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './outputs/grpo') + +# System prompt +SYSTEM_PROMPT = ( + 'You are an expert competition mathematician. ' + 'Solve the problem step by step. Put your final answer inside \\boxed{}. ' + 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.' +) + + +# ============================================================================ +# Reward +# ============================================================================ +class AoPSAccuracyReward(Reward): + """Accuracy reward via boxed answer extraction + robust equivalence matching.""" + + @staticmethod + def extract_boxed(text: str) -> str: + idx = text.rfind('\\boxed{') + if idx == -1: + return '' + start = idx + len('\\boxed{') + depth = 1 + j = start + while j < len(text) and depth > 0: + if text[j] == '{': + depth += 1 + elif text[j] == '}': + depth -= 1 + j += 1 + if depth == 0: + return text[start:j - 1].strip() + return '' + + _MCQ_GT_RE = re.compile( + r'^\\?(?:textbf|mathbf|text|mathrm)\{?\(?([A-E])[)}\s\\]*(.*)', + re.DOTALL) + _MCQ_PAREN_RE = re.compile(r'^\(?([A-E])\)?[.:\s\\]+(.*)', re.DOTALL) + _MCQ_SINGLE_LETTER_RE = re.compile(r'^[A-E]$') + _VAR_PREFIX_RE = re.compile( + r'^(?:[a-zA-Z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)', re.DOTALL) + _EQ_RHS_RE = re.compile(r'^.+=\s*(.+)$') + + @staticmethod + def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = ans.strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.replace(' ', '') + s = s.replace(r'\,', '') + s = s.replace(r'\;', '') + s = s.replace(r'\!', '') + s = re.sub(r'\\(?:text|mathrm|mathbf|textbf|operatorname)\{([^}]*)\}', r'\1', s) + s = s.replace(r'\displaystyle', '') + s = re.sub(r'\\(?:left|right)[.()\[\]|]', '', s) + s = s.replace(r'\dfrac', r'\frac') + s = s.replace(r'\tfrac', r'\frac') + # \frac shorthand without braces: \frac ab → \frac{a}{b} + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = s.strip('$').strip() + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\{\\circ\}|\^\\circ|°|\\circ', '', s) + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(m): + text = m.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start = pos + depth = 1 + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + denom = text[den_start:pos - 1] + return f'({numer})/({denom})' + + s = re.sub( + r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', + _frac_to_slash, s) + s = re.sub(r'(? str: + m = cls._VAR_PREFIX_RE.match(s) + return m.group(1).strip() if m else s + + @classmethod + def _extract_mcq_parts(cls, s: str): + m = cls._MCQ_GT_RE.match(s) + if m: + return m.group(1), m.group(2).strip() + m = cls._MCQ_PAREN_RE.match(s) + if m: + return m.group(1), m.group(2).strip() + m2 = re.search(r'\(?([A-E])\)?\s*$', s) + if m2 and len(s) > 3: + return m2.group(1), s[:m2.start()].strip() + return None, None + + @staticmethod + def _try_numeric_equal(a: str, b: str) -> bool: + import math + + def _try_eval(s: str): + try: + return float(s.replace('(', '').replace(')', '')) + except (ValueError, ZeroDivisionError): + pass + s_stripped = re.sub(r'[a-zA-Z]+$', '', s.replace('(', '').replace(')', '')).strip() + if s_stripped and s_stripped != s: + try: + return float(s_stripped) + except (ValueError, ZeroDivisionError): + pass + frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') + m = frac_re.match(s) + if m: + try: + return float(m.group(1)) / float(m.group(2)) + except (ValueError, ZeroDivisionError): + pass + expr = s + expr = expr.replace('\\pi', str(math.pi)) + expr = expr.replace('\\e', str(math.e)) + expr = re.sub(r'\\sqrt\{([^}]+)\}', r'(\1)**0.5', expr) + expr = re.sub(r'\\sqrt\[3\]\{([^}]+)\}', r'(\1)**(1/3)', expr) + expr = re.sub(r'\\sqrt\[([^]]+)\]\{([^}]+)\}', r'(\2)**(1/\1)', expr) + expr = expr.replace('{', '(').replace('}', ')') + expr = expr.replace('\\cdot', '*').replace('\\times', '*') + expr = re.sub(r'(\d)\(', r'\1*(', expr) + try: + val = eval(expr, {"__builtins__": {}, "math": math, "pi": math.pi, "e": math.e}, {}) + return float(val) + except Exception: + pass + return None + + va, vb = _try_eval(a), _try_eval(b) + if va is not None and vb is not None: + return abs(va - vb) < 1e-6 * max(1, abs(va), abs(vb)) + return False + + @classmethod + def _strip_quantifiers(cls, s: str) -> str: + """Strip universal/existential quantifier wrappers.""" + s = re.sub(r'^\\forall\s*\w+\s*\\in\s*\\mathbb\s*\{?[A-Z]\}?\s*[:,]\s*', '', s) + s = re.sub(r'\s*\(\\forall[^)]*\)\s*$', '', s) + s = re.sub(r'\s*\(for\s+all[^)]*\)\s*$', '', s, flags=re.IGNORECASE) + return s.strip() + + @classmethod + def _try_param_rename(cls, a: str, b: str) -> bool: + """Check if a == b up to consistent single free-parameter rename (ax vs cx).""" + if not a or not b or len(a) != len(b) or len(a) > 80: + return False + diffs = [(i, a[i], b[i]) for i in range(len(a)) if a[i] != b[i]] + if not diffs: + return False + src_chars = set(d[1] for d in diffs) + dst_chars = set(d[2] for d in diffs) + if len(src_chars) == 1 and len(dst_chars) == 1: + src, dst = src_chars.pop(), dst_chars.pop() + if src.isalpha() and dst.isalpha(): + return a.replace(src, dst) == b + return False + + @classmethod + def _normalize_tuple(cls, s: str) -> str: + # Strip set-builder conditions: \mid ... or | ... + s = re.sub(r'\\mid.*$', '', s) + s = re.sub(r'\|[^,]*$', '', s) + return re.sub(r'[\s()\[\]{}\\]', '', s) + + @classmethod + def _try_sympy_equal(cls, a: str, b: str) -> bool: + try: + from sympy.parsing.latex import parse_latex + from sympy import simplify, nsimplify + expr_a = parse_latex(a) + expr_b = parse_latex(b) + diff = simplify(nsimplify(expr_a - expr_b)) + return diff == 0 + except Exception: + return False + + @classmethod + def answers_match(cls, predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + + norm_p = cls.normalize_answer(predicted) + norm_r = cls.normalize_answer(reference) + + if norm_p == norm_r: + return True + if norm_p.lower() == norm_r.lower(): + return True + if cls._try_numeric_equal(norm_p, norm_r): + return True + + stripped_p = cls.normalize_answer(cls._strip_var_prefix(predicted)) + stripped_r = cls.normalize_answer(cls._strip_var_prefix(reference)) + if stripped_p and stripped_r and stripped_p == stripped_r: + return True + if stripped_p and stripped_r and cls._try_numeric_equal(stripped_p, stripped_r): + return True + + ref_letter, ref_value = cls._extract_mcq_parts(reference) + if ref_letter: + if norm_p == ref_letter or predicted.strip().upper() == ref_letter: + return True + if ref_value: + norm_ref_val = cls.normalize_answer(ref_value) + if norm_p == norm_ref_val or cls._try_numeric_equal(norm_p, norm_ref_val): + return True + pred_letter, pred_value = cls._extract_mcq_parts(predicted) + if pred_letter: + if norm_r == pred_letter or reference.strip().upper() == pred_letter: + return True + if pred_value: + norm_pred_val = cls.normalize_answer(pred_value) + if norm_r == norm_pred_val or cls._try_numeric_equal(norm_r, norm_pred_val): + return True + if cls._MCQ_SINGLE_LETTER_RE.match(reference.strip()): + if cls._MCQ_SINGLE_LETTER_RE.match(predicted.strip().upper()): + return predicted.strip().upper() == reference.strip().upper() + + tuple_p = cls._normalize_tuple(norm_p) + tuple_r = cls._normalize_tuple(norm_r) + if ',' in tuple_p and tuple_p == tuple_r: + return True + if stripped_p and stripped_r: + tuple_sp = cls._normalize_tuple(stripped_p) + tuple_sr = cls._normalize_tuple(stripped_r) + if ',' in tuple_sp and tuple_sp == tuple_sr: + return True + + if '=' in norm_r and '=' not in norm_p: + parts = norm_r.split('=') + for part in parts: + part = part.strip() + if part == norm_p or cls._try_numeric_equal(part, norm_p): + return True + if '=' in norm_p and '=' not in norm_r: + parts = norm_p.split('=') + for part in parts: + part = part.strip() + if part == norm_r or cls._try_numeric_equal(part, norm_r): + return True + if '=' in norm_p and '=' in norm_r: + pp = [x.strip() for x in norm_p.split('=')] + rp = [x.strip() for x in norm_r.split('=')] + if set(pp) == set(rp): + return True + + def _sort_factors(s): + tokens = re.findall(r'\\?[a-zA-Z]+\{[^}]*\}|\\?[a-zA-Z]+|\d+|[^a-zA-Z\d\\{}]', s) + return ''.join(sorted(tokens)) + if _sort_factors(norm_p) == _sort_factors(norm_r): + return True + + if cls._try_sympy_equal(predicted, reference): + return True + + # --- Strategy 9b: quantifier stripping + param rename --- + q_stripped_r = cls._strip_quantifiers(reference) + q_stripped_p = cls._strip_quantifiers(predicted) + if q_stripped_r != reference or q_stripped_p != predicted: + norm_qr = cls.normalize_answer(cls._strip_var_prefix(q_stripped_r)) + norm_qp = cls.normalize_answer(cls._strip_var_prefix(q_stripped_p)) + if norm_qr and norm_qp: + if norm_qr == norm_qp: + return True + if cls._try_param_rename(norm_qr, norm_qp): + return True + + # --- Strategy 10: param rename on var-prefix-stripped forms --- + if stripped_p and stripped_r and cls._try_param_rename(stripped_p, stripped_r): + return True + + return False + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + user_data = traj.get('user_data') or [] + gt = '' + for item in user_data: + if item[0] == 'ground_truth': + gt = item[1] + break + predicted = self.extract_boxed(completion) + correct = self.answers_match(predicted, gt) + rewards.append(1.0 if correct else 0.0) + return rewards + + +class FormatReward(Reward): + """Reward for having \\boxed{} in the output.""" + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + has_boxed = '\\boxed{' in completion + rewards.append(0.5 if has_boxed else 0.0) + return rewards + + +class GibberishPenalty(Reward): + """Negative reward for degenerate outputs (gibberish/random unicode tail).""" + + TAIL_CHARS = 400 + GIBBERISH_THRESHOLD = 0.20 + + @classmethod + def is_gibberish(cls, text: str) -> bool: + if not text: + return False + tail = text[-cls.TAIL_CHARS:] if len(text) > cls.TAIL_CHARS else text + non_math_non_ascii = 0 + for c in tail: + code = ord(c) + if code > 127 and not (0x4e00 <= code <= 0x9fff): + non_math_non_ascii += 1 + return non_math_non_ascii > len(tail) * cls.GIBBERISH_THRESHOLD + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + rewards.append(-0.5 if self.is_gibberish(completion) else 0.0) + return rewards + + +def compute_rewards(trajectories: List[Dict[str, Any]] + ) -> Tuple[List[float], List[float], List[float]]: + acc_fn = AoPSAccuracyReward() + fmt_fn = FormatReward() + gib_fn = GibberishPenalty() + acc = acc_fn(trajectories) + fmt = fmt_fn(trajectories) + gib = gib_fn(trajectories) + total = [a + f + g for a, f, g in zip(acc, fmt, gib)] + return total, fmt, acc + + +# ============================================================================ +# Dataset: AoPS boxed problems +# ============================================================================ +def create_aops_dataset(): + """Load AoPS and create GRPO-style dataset (prompt only, with ground_truth in user_data).""" + from modelscope import MsDataset + from twinkle.data_format import Message, Trajectory + + ds = MsDataset.load(AOPS_DATASET_ID, split='train', + download_mode='reuse_dataset_if_exists') + rows = [] + for row in ds: + if not row['metadata'].get('boxed'): + continue + ref = AoPSAccuracyReward.extract_boxed(row['solution']) + if not ref: + continue + rows.append({'problem': row['problem'], 'ground_truth': ref}) + + logger.info(f'[aops] loaded {len(rows)} boxed problems') + rng = random.Random(AOPS_SEED) + rng.shuffle(rows) + + trajectories = [] + for r in rows: + traj = Trajectory( + messages=[ + Message(role='system', content=SYSTEM_PROMPT), + Message(role='user', content=r['problem']), + ], + user_data=[('ground_truth', r['ground_truth'])], + ) + trajectories.append(traj) + + data_meta = DatasetMeta(data=trajectories) + dataset = Dataset(data_meta) + dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, + max_length=16384, truncation_strategy='delete', + enable_thinking=True) + dataset.encode(add_generation_prompt=True) + return dataset + + +# ============================================================================ +# Main +# ============================================================================ +def main(): + sampler_start = 0 + model_start = sampler_start + SAMPLER_GPUS + + device_groups = [ + DeviceGroup(name='sampler', ranks=list(range(sampler_start, model_start)), + device_type='GPU'), + DeviceGroup(name='model', ranks=list(range(model_start, NUM_GPUS)), + device_type='GPU'), + ] + + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=2) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, + groups=device_groups, lazy_collect=False) + + # -- Training model (full-parameter) -- + model = TransformersModel( + model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') + model.set_optimizer('AdamW', lr=LEARNING_RATE) + model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + model.set_loss('GSPOLoss', epsilon=0.2, epsilon_high=0.28, beta=0.04) + model.set_processor(InputProcessor) + model.set_template('Qwen3_5Template', model_id=MODEL_ID, + enable_thinking=True, max_length=32768) + + # -- Rollout sampler -- + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 32768, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, + enable_thinking=True, max_length=32768) + + # -- Checkpoint & DataLoader -- + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + + GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS + dataloader = DataLoader( + dataset=create_aops_dataset, + batch_size=GLOBAL_BATCH_SIZE, + min_batch_size=GLOBAL_BATCH_SIZE, + device_mesh=model_mesh, + remote_group='model', + ) + + advantage_fn = GRPOAdvantage() + metrics = CompletionRewardMetric() + sampling_params = SamplingParams( + max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, + temperature=1.0, top_p=0.95) + + optim_step = 0 + logger.info('Starting pure GRPO training (no RAG)') + logger.info(get_device_placement()) + + # -- Diagnostics -- + os.makedirs(OUTPUT_DIR, exist_ok=True) + diag_path = os.path.join(OUTPUT_DIR, 'diagnostics.jsonl') + diag_f = open(diag_path, 'w', encoding='utf-8') + logger.info(f'[diag] diagnostics → {diag_path}') + + def _content_to_str(content): + if isinstance(content, str): + return content + if isinstance(content, list): + return ''.join( + b.get('text', '') if isinstance(b, dict) else str(b) + for b in content) + return str(content) + + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + + metrics.reset() + + # Build prompts (direct, no RAG) + prompts = [] + for item in batch: + msgs = item.get('messages', []) + prob = '' + for m in msgs: + if m.get('role') == 'user': + prob = m.get('content', '') + if isinstance(prob, list): + prob = ''.join(p.get('text', '') for p in prob if isinstance(p, dict)) + break + ud = item.get('user_data', []) + gt = '' + for pair in ud: + if pair[0] == 'ground_truth': + gt = pair[1] + break + prompts.append({ + 'messages': [ + {'role': 'system', 'content': SYSTEM_PROMPT}, + {'role': 'user', 'content': prob}, + ], + 'user_data': [('ground_truth', gt)], + }) + + # Expand for NUM_GENERATIONS and sample + expand_prompts = [] + for prompt in prompts: + expand_prompts.extend([prompt] * NUM_GENERATIONS) + + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + + sample_responses = sampler.sample(expand_prompts, sampling_params) + + # Collect rollouts + all_input_data: List[Dict[str, Any]] = [] + all_old_logps: List[List[float]] = [] + all_completion_lengths: List[int] = [] + + for sample_response in sample_responses: + for sequence in sample_response.sequences: + all_input_data.append(sequence.new_input_feature) + all_old_logps.append([logprob[0][1] for logprob in sequence.logprobs]) + all_completion_lengths.append(len(sequence.tokens)) + + # Rewards + total_rewards, format_rewards, accuracy_rewards = compute_rewards(all_input_data) + + # Zero out rewards for rollouts that hit the max_tokens ceiling + max_len_threshold = int(MAX_NEW_TOKENS * 0.95) + for i in range(len(all_input_data)): + if all_completion_lengths[i] >= max_len_threshold: + total_rewards[i] = 0.0 + accuracy_rewards[i] = 0.0 + format_rewards[i] = 0.0 + + # Per-step reward summary + n_correct = sum(1 for a in accuracy_rewards if a > 0) + diag_f.write(json.dumps({ + 'step': optim_step, 'type': 'reward_summary', + 'n_samples': len(accuracy_rewards), + 'accuracy': n_correct / len(accuracy_rewards) if accuracy_rewards else 0, + 'mean_reward': sum(total_rewards) / len(total_rewards) if total_rewards else 0, + }, ensure_ascii=False) + '\n') + + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={ + 'total': total_rewards, + 'format': format_rewards, + 'accuracy': accuracy_rewards, + }, + ) + + # GRPO advantage + advantages = advantage_fn( + total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + if ADV_CLIP > 0: + advantages = [max(-ADV_CLIP, min(ADV_CLIP, a)) for a in advantages] + + # Log rollout responses + _extract_boxed = AoPSAccuracyReward.extract_boxed + for ridx, traj in enumerate(all_input_data): + msgs = traj.get('messages', []) + assistant_text = _content_to_str(next( + (m['content'] for m in reversed(msgs) if m.get('role') == 'assistant'), '')) + user_text = _content_to_str(next( + (m['content'] for m in msgs if m.get('role') == 'user'), '')) + user_data = traj.get('user_data') or [] + gt = next((v for k, v in user_data if k == 'ground_truth'), '') + problem_idx = ridx // NUM_GENERATIONS + grp_start = problem_idx * NUM_GENERATIONS + grp_end = grp_start + NUM_GENERATIONS + grp_acc = sum(accuracy_rewards[grp_start:grp_end]) / NUM_GENERATIONS + + diag_f.write(json.dumps({ + 'step': optim_step, 'type': 'rollout', + 'idx': ridx, + 'problem_idx': problem_idx, + 'problem': user_text, + 'response': assistant_text, + 'ground_truth': gt, + 'predicted': _extract_boxed(assistant_text), + 'reward': total_rewards[ridx], + 'accuracy_reward': accuracy_rewards[ridx], + 'format_reward': format_rewards[ridx], + 'advantage': advantages[ridx], + 'completion_length': all_completion_lengths[ridx], + 'group_accuracy': grp_acc, + }, ensure_ascii=False) + '\n') + + diag_f.flush() + + # Filter out low-signal problem groups (DAPO-style dynamic sampling) + # Skip groups where accuracy is too low (<0.1) or too high (>0.9) + # to avoid gradient dominated by gibberish/format noise or no learning signal. + filtered_inputs, filtered_old_logps, filtered_advantages = [], [], [] + for g in range(BATCH_SIZE): + g_start = g * NUM_GENERATIONS + g_end = g_start + NUM_GENERATIONS + grp_adv = advantages[g_start:g_end] + if all(abs(a) < 1e-8 for a in grp_adv): + continue + grp_acc_rate = sum(accuracy_rewards[g_start:g_end]) / NUM_GENERATIONS + if grp_acc_rate < 0.2 or grp_acc_rate > 0.8: + continue + filtered_inputs.extend(all_input_data[g_start:g_end]) + filtered_old_logps.extend(all_old_logps[g_start:g_end]) + filtered_advantages.extend(grp_adv) + + # Mini-batch training with gradient accumulation + # Process MICRO_BATCH_SIZE samples per forward, accumulate grad_accum_steps + # times before one optimizer step. clip_grad_norm normalizes by accumulated + # num_tokens, ensuring mathematical equivalence with larger batch forward. + total_completions = len(filtered_inputs) + if total_completions == 0: + logger.info(f'[Step {optim_step}] all groups filtered (uniform rewards), skip training') + continue + + grad_accum_steps = MINI_BATCH_SIZE // MICRO_BATCH_SIZE + accum_count = 0 + for mb_start in range(0, total_completions, MICRO_BATCH_SIZE): + mb_end = min(mb_start + MICRO_BATCH_SIZE, total_completions) + mb_inputs = filtered_inputs[mb_start:mb_end] + mb_old_logps = filtered_old_logps[mb_start:mb_end] + mb_advantages = filtered_advantages[mb_start:mb_end] + + outputs = model.forward_backward( + inputs=mb_inputs, + old_logps=mb_old_logps, + ref_logps=mb_old_logps, + advantages=mb_advantages, + ) + accum_count += 1 + + if accum_count % grad_accum_steps == 0: + skip_step = False + try: + loss_val = outputs.get('loss', None) + if loss_val is not None: + if hasattr(loss_val, 'item'): + loss_val = loss_val.item() + if loss_val > LOSS_SPIKE_THRESHOLD: + skip_step = True + logger.warning( + f'[Step {optim_step}] Loss spike: {loss_val:.4f} > ' + f'{LOSS_SPIKE_THRESHOLD}, skipping update') + except Exception: + pass + + if skip_step: + model.zero_grad() + else: + model.clip_grad_and_step() + optim_step += 1 + + if optim_step >= MAX_STEPS: + break + if optim_step % SAVE_STEPS == 0: + model.save(f'grpo-checkpoint-{optim_step}') + + # Flush remaining accumulated gradients (incomplete window at tail) + if accum_count % grad_accum_steps != 0: + skip_step = False + try: + loss_val = outputs.get('loss', None) + if loss_val is not None: + if hasattr(loss_val, 'item'): + loss_val = loss_val.item() + if loss_val > LOSS_SPIKE_THRESHOLD: + skip_step = True + logger.warning( + f'[Step {optim_step}] Loss spike (tail): {loss_val:.4f} > ' + f'{LOSS_SPIKE_THRESHOLD}, skipping update') + except Exception: + pass + + if skip_step: + model.zero_grad() + else: + model.clip_grad_and_step() + optim_step += 1 + + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True)) + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + + diag_f.close() + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('grpo-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py new file mode 100644 index 00000000..b35b3706 --- /dev/null +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -0,0 +1,1480 @@ +"""RAG-hint GRPO training: retrieve thinking traces and condense as hints for RL. + +Architecture (8 GPUs): + - 1 GPU: condenser (vLLM, compress retrieved traces) + - 1 GPU: embedding model (encode queries for retrieval) + - 4 GPUs: sampler/rollout (vLLM TP=4) + - 2 GPUs: training model (FSDP/DP) + +Pipeline per step: + 1. DataLoader yields a batch of math problems + 2. [Async] Embedding model encodes problems → retrieve from LanceDB → condenser compresses + 3. Build RAG-hint prompts (one-shot in system, with analysis prefix) + 4. Sampler generates rollouts (response starts with forced analysis prefix) + 5. Reward (accuracy + format) → GRPO advantage → model update + +Launch: + python cookbook/exp/rl/rag_hint_grpo.py +""" +import json +import os +import re +import random +import threading +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import torch + +import twinkle +from twinkle import DeviceMesh, DeviceGroup, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.loss import InfonceLoss +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.reward.base import Reward +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient + +logger = get_logger() + +# ============================================================================ +# Configuration +# ============================================================================ +MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') + +# GPU layout: 1 condenser + 1 embedding + 4 rollout + 2 train = 8 +CONDENSER_GPUS = int(os.environ.get('CONDENSER_GPUS', 1)) +EMB_GPUS = int(os.environ.get('EMB_GPUS', 1)) +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 2)) +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) +NUM_GPUS = CONDENSER_GPUS + EMB_GPUS + SAMPLER_GPUS + MODEL_GPUS + +# Training hyperparams +NUM_GENERATIONS = int(os.environ.get('NUM_GENERATIONS', 8)) +MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 32768)) +LEARNING_RATE = float(os.environ.get('LR', 1e-5)) +MAX_STEPS = int(os.environ.get('MAX_STEPS', 5000)) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 8)) +MINI_BATCH_SIZE = int(os.environ.get('MINI_BATCH_SIZE', 8)) +MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 8)) +GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 1)) +SAVE_STEPS = int(os.environ.get('SAVE_STEPS', 100)) +ADV_CLIP = float(os.environ.get('ADV_CLIP', 1.0)) +LOSS_SPIKE_THRESHOLD = float(os.environ.get('LOSS_SPIKE_THRESHOLD', 10.0)) + +# RAG config +DB_PATH = os.environ.get('DB_PATH', './output.oldemb/thinking_rag/lance.db') +DB_TABLE = os.environ.get('DB_TABLE', 'thinking_traces') +TOP_K = int(os.environ.get('TOP_K', 2)) +SIM_THRESHOLD = float(os.environ.get('SIM_THRESHOLD', 0.75)) +MAX_TRACE_LEN = int(os.environ.get('MAX_TRACE_LEN', 8192)) +EMBED_MODEL_ID = os.environ.get( + 'EMBED_MODEL_ID', 'output.oldemb/embedding_full_transformers/last-checkpoint') +EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 32000)) + +# Condenser config +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +CONDENSE_API_KEY = os.environ.get('COMPRESS_API_KEY', '') +CONDENSE_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') +CONDENSE_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') +CONDENSE_API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 16)) +CONDENSE_TEMPERATURE = 0.2 +CONDENSE_MAX_TOKENS = 8192 + +# Dataset +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +AOPS_SEED = int(os.environ.get('AOPS_SEED', 100)) + +# Decontamination & RAG fallback +DECONTAM_THRESHOLD = float(os.environ.get('DECONTAM_THRESHOLD', 0.20)) +RAG_FALLBACK_SIM = float(os.environ.get('RAG_FALLBACK_SIM', 0.60)) + +# Output / diagnostics +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './outputs/rag_hint_grpo') + +# Forced analysis prefix appended at the start of assistant response +ANALYSIS_PREFIX = '' + +# Fixed opening inside block — model must produce this EXACT prefix +HINT_REQUIRED_PREFIX = "Let's analyze the RAG example step by step." + +# Hint analysis config (API pre-analysis) +HINT_ANALYSIS_MAX_TOKENS = int(os.environ.get('HINT_ANALYSIS_MAX_TOKENS', 400)) +HINT_ANALYSIS_TEMPERATURE = 0.3 + +# ============================================================================ +# Condenser prompt (strategy-level extraction) +# ============================================================================ +COMPRESS_SYSTEM = """\ +You are a reasoning-trace condenser. Given a verbose reasoning trace, \ +extract the TRANSFERABLE KNOWLEDGE as an EXECUTABLE SOLUTION SKELETON \ +that would help a reader solve SIMILAR problems in the same domain. + +The reader will apply this knowledge to a DIFFERENT problem, so focus on what transfers. \ +NEVER output the final answer or conclusion of the original problem. \ +NEVER include problem-specific numeric results. + +Principles: +1. OUTPUT AN EXECUTABLE STEP CHAIN: numbered steps that a solver can directly follow. \ +Each step should state WHAT/WHY/HOW (with the formula/technique), not just name the concept. +2. INCLUDE FULL FORMULAS: theorems, identities — state each with COMPLETE MATHEMATICAL EXPRESSION. +3. STATE APPLICABILITY: what structural features signal that this approach works. +4. PRESERVE KEY INSIGHTS: non-obvious ideas that make the approach work. +5. REMOVE: problem-specific numeric calculations, final answers, dead-end explorations, hesitations. +6. FORMAT: Start with "Applicability:" one-line, then numbered steps. Keep concise. +7. NO meta-commentary. NO preamble. NO final answer. +""" + +COMPRESS_USER = ( + '## Reader Problem (context only — do NOT solve it)\n{query}\n\n' + '## Reasoning Trace to Condense\n{text}') + +# ============================================================================ +# RAG system prompt template (few-shot in system) +# ============================================================================ +SYSTEM_WITH_RAG_HEADER = ( + 'You are an expert competition mathematician. ' + 'Below are condensed reasoning examples from similar problems.\n\n' + '## Output Format (STRICT)\n' + 'Your response MUST begin with a block as the VERY FIRST content. ' + 'Do NOT output any text before .\n\n' + 'The block MUST start with EXACTLY this sentence (copy verbatim):\n' + '"Let\'s analyze the RAG example step by step."\n\n' + 'Then continue your analysis:\n' + '- Walk through each example\'s methodology and identify which steps, ' + 'formulas, and concepts are CORRECT and APPLICABLE to the current problem.\n' + '- Identify which parts are WRONG, IRRELEVANT, MISLEADING, or based on ' + 'assumptions that do NOT hold for this problem.\n' + '- End with a one-line verdict: "Useful: ..." and "Discard: ..."\n\n' + 'Example format:\n' + '\n' + "Let's analyze the RAG example step by step.\n" + '- Example 1: The ansatz f(x)=x^n is APPLICABLE because ... However, ' + 'the uniqueness argument via continuity is UNNECESSARY for this problem.\n' + '- Useful: power function ansatz, linear combination check.\n' + '- Discard: continuity assumption, specific numeric result.\n' + '\n\n' + 'After the block, solve the actual problem step by step using ONLY ' + 'the validated useful parts. Put your final answer inside \\boxed{}. ' + 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.\n\n' +) + +# System prompt for pre-analyzed RAG (hint analysis done by API, model just solves) +_PREANALYSIS_BEFORE = ( + 'You are an expert competition mathematician.\n\n' + '## RAG Analysis (pre-computed)\n' +) +_PREANALYSIS_AFTER = ( + '\n\n## Instructions\n' + 'Use the useful methods/formulas identified above to solve the problem. ' + 'Ignore anything marked as irrelevant. ' + 'Solve step by step and put your final answer inside \\boxed{}. ' + 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.' +) + + +def build_preanalysis_system(hint_analysis: str) -> str: + """Build system prompt with pre-analyzed hint. Uses concatenation to avoid .format() issues with math braces.""" + return _PREANALYSIS_BEFORE + hint_analysis + _PREANALYSIS_AFTER + +# API prompt for hint analysis generation +HINT_ANALYSIS_SYSTEM = ( + 'You are a mathematical methodology analyst. ' + 'Given a target problem and a condensed reasoning trace from a SIMILAR (but different) problem, ' + 'analyze which methods, formulas, and techniques from the trace are APPLICABLE to the target problem ' + 'and which are IRRELEVANT or MISLEADING.\n\n' + 'Output format (strict):\n' + '- Useful: [list specific methods/formulas/techniques that transfer to the target]\n' + '- Discard: [list parts that are irrelevant or would mislead]\n' + '- Key insight: [one sentence on how to apply the useful parts]\n\n' + 'Rules:\n' + '1. Be concise — at most 200 words total.\n' + '2. Focus ONLY on transferable methodology, never solve the target problem.\n' + '3. Never output the answer to either problem.\n' + '4. If the trace is entirely irrelevant, say "Useful: None. Discard: All."' +) + +HINT_ANALYSIS_USER = ( + '## Target Problem\n{query}\n\n' + '## Condensed Trace (from similar problem)\n{thinking}' +) + +EXAMPLE_TEMPLATE = ( + '--- Example {idx} ---\n' + 'Problem: {example_query}\n' + 'Methodology:\n{example_thinking}\n' + '--- End Example {idx} ---\n' +) + +SYSTEM_DIRECT = ( + 'You are an expert competition mathematician. ' + 'Solve the problem step by step. Put your final answer inside \\boxed{}. ' + 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.' +) + + +# ============================================================================ +# Condenser utilities +# ============================================================================ +_api_semaphore = threading.Semaphore(CONDENSE_API_CONCURRENCY) + + +def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Optional[str]: + _api_semaphore.acquire() + try: + trajectory = {'messages': messages} + sp = SamplingParams(temperature=CONDENSE_TEMPERATURE, max_tokens=CONDENSE_MAX_TOKENS) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + content = (reply.get('content') or '').strip() + if not content: + return None + return content + except Exception as exc: + logger.warning(f'[condense-api] error: {exc}') + return None + finally: + _api_semaphore.release() + + +def _api_hint_analysis_batch( + api_client: OpenAIClient, + problems: List[str], + condensed_examples: List[List[Dict[str, str]]], +) -> List[Optional[str]]: + """Call API to pre-analyze RAG relevance for each problem. ~300 tokens per call.""" + results: List[Optional[str]] = [None] * len(problems) + tasks = [] # (idx, messages) + for i, prob in enumerate(problems): + if not condensed_examples[i]: + continue + # Merge all condensed traces into one block + traces = [] + for ex in condensed_examples[i]: + traces.append(ex.get('thinking', '')) + merged_thinking = '\n---\n'.join(traces) + user_msg = HINT_ANALYSIS_USER.format(query=prob, thinking=merged_thinking) + msgs = [ + {'role': 'system', 'content': HINT_ANALYSIS_SYSTEM}, + {'role': 'user', 'content': user_msg}, + ] + tasks.append((i, msgs)) + + if not tasks: + return results + + def _call_one(idx, msgs): + _api_semaphore.acquire() + try: + trajectory = {'messages': msgs} + sp = SamplingParams( + temperature=HINT_ANALYSIS_TEMPERATURE, + max_tokens=HINT_ANALYSIS_MAX_TOKENS) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + content = (reply.get('content') or '').strip() + return idx, content if content else None + except Exception as exc: + logger.warning(f'[hint-analysis] error for idx={idx}: {exc}') + return idx, None + finally: + _api_semaphore.release() + + with ThreadPoolExecutor(max_workers=min(len(tasks), CONDENSE_API_CONCURRENCY)) as pool: + futs = [pool.submit(_call_one, idx, msgs) for idx, msgs in tasks] + for fut in as_completed(futs): + idx, analysis = fut.result() + results[idx] = analysis + + n_success = sum(1 for r in results if r) + logger.info(f'[hint-analysis] completed {n_success}/{len(tasks)} analyses') + return results + + +# ============================================================================ +# Embedding & Retrieval +# ============================================================================ +def _normalize_for_ngram(text: str) -> str: + """Normalize text for n-gram comparison: strip LaTeX markup, lowercase.""" + text = text.lower() + text = re.sub(r'\$+', '', text) + text = re.sub(r'\\[a-z]+\{([^}]*)\}', r'\1', text) + text = re.sub(r'\\[a-z]+', ' ', text) + text = re.sub(r'[{}\\^_$]', '', text) + text = re.sub(r'\s+', ' ', text).strip() + return text + + +def _ngram_jaccard(text_a: str, text_b: str, n: int = 13) -> float: + """13-gram character-level Jaccard similarity for decontamination.""" + a = _normalize_for_ngram(text_a) + b = _normalize_for_ngram(text_b) + if len(a) < n or len(b) < n: + return 0.0 + grams_a = set(a[i:i + n] for i in range(len(a) - n + 1)) + grams_b = set(b[i:i + n] for i in range(len(b) - n + 1)) + if not grams_a or not grams_b: + return 0.0 + return len(grams_a & grams_b) / len(grams_a | grams_b) + + +def _wrap_anchor(text: str) -> List[Dict[str, str]]: + return [ + {'role': 'user', 'content': text}, + {'role': 'assistant', 'content': 'Match the correct response here.'}, + ] + + +_DECONTAM_JUDGE_PROMPT = ( + 'We are building a RAG-augmented math training system. Problem A is the test ' + 'question; Problem B was retrieved from a knowledge base.\n' + 'Answer YES only if A and B are essentially the SAME specific problem — ' + 'i.e. solving B directly gives you A\'s answer (just different wording/notation/' + 'format/negation).\n' + 'Answer NO if they merely share the same method/topic but have different ' + 'specific values, equations, or geometric configurations — learning B\'s ' + 'approach still requires independent work to solve A.\n' + 'Problem A: {prob_a}\n' + 'Problem B: {prob_b}\n' + 'Answer only YES or NO.' +) + + +def _llm_judge_same_problem( + api_client, pairs: List[tuple], +) -> List[bool]: + """Batch LLM judge: are (problem_a, problem_b) the same problem? + + Each pair text is truncated to 200 chars to keep latency low. + Returns list of bools (True = same problem = should filter). + """ + if not pairs or not api_client: + return [False] * len(pairs) + + results = [False] * len(pairs) + + def _judge_one(idx, pa, pb): + prompt = _DECONTAM_JUDGE_PROMPT.format(prob_a=pa, prob_b=pb) + msgs = [{'role': 'user', 'content': prompt}] + try: + trajectory = {'messages': msgs} + sp = SamplingParams(temperature=0.1, max_tokens=8) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + answer = (reply.get('content') or '').strip().upper() + return idx, 'YES' in answer + except Exception: + return idx, False + + with ThreadPoolExecutor(max_workers=min(len(pairs), CONDENSE_API_CONCURRENCY)) as pool: + futs = [pool.submit(_judge_one, i, pa, pb) for i, (pa, pb) in enumerate(pairs)] + for fut in as_completed(futs): + idx, is_same = fut.result() + results[idx] = is_same + return results + + +def get_embeddings(model: TransformersModel, template: Qwen3_5Template, + texts: List[str], dp_size: int) -> np.ndarray: + if not texts: + return np.zeros((0,), dtype=np.float32) + n = len(texts) + pad_n = (-n) % dp_size + padded = list(texts) + [' '] * pad_n if pad_n else list(texts) + features = [] + for t in padded: + feat = template.encode({'messages': _wrap_anchor(t or ' ')}) + feat['labels'] = [1] + features.append(feat) + out = model.forward_only(inputs=features, task='embedding', return_logits=True) + emb = out['embeddings'] + if isinstance(emb, torch.Tensor): + emb = emb.detach().to(torch.float32).cpu().numpy() + emb = np.asarray(emb, dtype=np.float32) + return emb[:n] if pad_n else emb + + +def retrieve_topk(tbl, query_vecs: np.ndarray, problems: List[str], + sim_threshold: float + ) -> List[List[Dict[str, Any]]]: + """Retrieve top-K thinking_raw per query with decontamination and length filter. + + Returns per-query list of dicts with keys: query, thinking, sim. + """ + results = [] + decontam_skipped = 0 + for qi, vec in enumerate(query_vecs): + hits = ( + tbl.search(vec.astype(np.float32).tolist()) + .metric('dot') + .limit(TOP_K + 50) + .select(['query_raw', 'thinking_raw', '_distance']) + .to_list() + ) + matched = [] + problem_text = problems[qi] if problems else '' + for h in hits: + if len(matched) >= TOP_K: + break + sim = 1.0 - h.get('_distance', 0.0) + if sim < sim_threshold: + continue + q = h.get('query_raw', '') + t = h.get('thinking_raw', '') + if not t: + continue + # Decontamination: skip if retrieved problem is too similar to current + if DECONTAM_THRESHOLD > 0 and problem_text and q: + if _ngram_jaccard(problem_text, q) > DECONTAM_THRESHOLD: + decontam_skipped += 1 + continue + # Drop traces exceeding max length (don't truncate — they'll be condensed poorly) + if len(t) > MAX_TRACE_LEN * 4: + continue + matched.append({'query': q, 'thinking': t, 'sim': sim}) + results.append(matched) + if decontam_skipped > 0: + logger.info(f'[decontam] skipped {decontam_skipped} leaked retrievals') + return results + + +# ============================================================================ +# Reward +# ============================================================================ +class AoPSAccuracyReward(Reward): + """Accuracy reward via boxed answer extraction + robust equivalence matching.""" + + @staticmethod + def extract_boxed(text: str) -> str: + idx = text.rfind('\\boxed{') + if idx == -1: + return '' + start = idx + len('\\boxed{') + depth = 1 + j = start + while j < len(text) and depth > 0: + if text[j] == '{': + depth += 1 + elif text[j] == '}': + depth -= 1 + j += 1 + if depth == 0: + return text[start:j - 1].strip() + return '' + + # --- MCQ letter regex (matches \textbf{(C) }value, (C) value, etc.) --- + _MCQ_GT_RE = re.compile( + r'^\\?(?:textbf|mathbf|text|mathrm)\{?\(?([A-E])[)}\s\\]*(.*)', + re.DOTALL) + _MCQ_PAREN_RE = re.compile(r'^\(?([A-E])\)?[.:\s\\]+(.*)', re.DOTALL) + _MCQ_SINGLE_LETTER_RE = re.compile(r'^[A-E]$') + # variable prefix: f(x)=..., m=..., N=..., P(n+1)=..., (x,y)=... + _VAR_PREFIX_RE = re.compile( + r'^(?:[a-zA-Z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)', re.DOTALL) + # GT with derivation: "18×1+999×2=2016" → extract RHS + _EQ_RHS_RE = re.compile(r'^.+=\s*(.+)$') + + @staticmethod + def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = ans.strip() + # Pure MCQ letter + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.replace(' ', '') + s = s.replace(r'\,', '') + s = s.replace(r'\;', '') + s = s.replace(r'\!', '') + # Remove text-mode wrappers but keep content (unwrap braces) + s = re.sub(r'\\(?:text|mathrm|mathbf|textbf|operatorname)\{([^}]*)\}', r'\1', s) + s = s.replace(r'\displaystyle', '') + s = re.sub(r'\\(?:left|right)[.()\[\]|]', '', s) + s = s.replace(r'\dfrac', r'\frac') + s = s.replace(r'\tfrac', r'\frac') + # \frac shorthand without braces: \frac ab → \frac{a}{b} + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = s.strip('$').strip() + # Remove trailing unit braces: {cm}, {kg}, etc. + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + # Degree normalization — strip entirely (degrees are contextual) + s = re.sub(r'\^\{\\circ\}|\^\\circ|°|\\circ', '', s) + # Remove \quad, \qquad, \ etc spacing + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + # Normalize minus: \minus{} → - + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(m): + text = m.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start = pos + depth = 1 + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + denom = text[den_start:pos - 1] + return f'({numer})/({denom})' + + s = re.sub( + r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', + _frac_to_slash, s) + s = re.sub(r'(? str: + """Strip variable assignment prefix: 'f(x)=x+1' → 'x+1', 'N=1006' → '1006'.""" + m = cls._VAR_PREFIX_RE.match(s) + return m.group(1).strip() if m else s + + @classmethod + def _extract_mcq_parts(cls, s: str): + """Extract (letter, value) from MCQ-formatted string. Returns (None, None) if not MCQ.""" + m = cls._MCQ_GT_RE.match(s) + if m: + return m.group(1), m.group(2).strip() + m = cls._MCQ_PAREN_RE.match(s) + if m: + return m.group(1), m.group(2).strip() + # GT ends with " (A)" pattern: "2+2\sqrt{7} (A)" + m2 = re.search(r'\(?([A-E])\)?\s*$', s) + if m2 and len(s) > 3: + return m2.group(1), s[:m2.start()].strip() + return None, None + + @staticmethod + def _try_numeric_equal(a: str, b: str) -> bool: + """Try numeric equality after normalization. Handles fracs and simple expressions.""" + import math + + def _try_eval(s: str): + # Direct float + try: + return float(s.replace('(', '').replace(')', '')) + except (ValueError, ZeroDivisionError): + pass + # Strip trailing unit-like suffix and retry + s_stripped = re.sub(r'[a-zA-Z]+$', '', s.replace('(', '').replace(')', '')).strip() + if s_stripped and s_stripped != s: + try: + return float(s_stripped) + except (ValueError, ZeroDivisionError): + pass + # Fraction pattern (a)/(b) + frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') + m = frac_re.match(s) + if m: + try: + return float(m.group(1)) / float(m.group(2)) + except (ValueError, ZeroDivisionError): + pass + # Try evaluating simple math expressions (pi, sqrt, etc.) + expr = s + expr = expr.replace('\\pi', str(math.pi)) + expr = expr.replace('\\e', str(math.e)) + expr = re.sub(r'\\sqrt\{([^}]+)\}', r'(\1)**0.5', expr) + expr = re.sub(r'\\sqrt\[3\]\{([^}]+)\}', r'(\1)**(1/3)', expr) + expr = re.sub(r'\\sqrt\[([^]]+)\]\{([^}]+)\}', r'(\2)**(1/\1)', expr) + expr = expr.replace('{', '(').replace('}', ')') + expr = expr.replace('\\cdot', '*').replace('\\times', '*') + expr = re.sub(r'(\d)\(', r'\1*(', expr) + try: + val = eval(expr, {"__builtins__": {}, "math": math, "pi": math.pi, "e": math.e}, {}) + return float(val) + except Exception: + pass + return None + + va, vb = _try_eval(a), _try_eval(b) + if va is not None and vb is not None: + return abs(va - vb) < 1e-6 * max(1, abs(va), abs(vb)) + return False + + @classmethod + def _strip_quantifiers(cls, s: str) -> str: + """Strip universal/existential quantifier wrappers.""" + s = re.sub(r'^\\forall\s*\w+\s*\\in\s*\\mathbb\s*\{?[A-Z]\}?\s*[:,]\s*', '', s) + s = re.sub(r'\s*\(\\forall[^)]*\)\s*$', '', s) + s = re.sub(r'\s*\(for\s+all[^)]*\)\s*$', '', s, flags=re.IGNORECASE) + return s.strip() + + @classmethod + def _try_param_rename(cls, a: str, b: str) -> bool: + """Check if a == b up to consistent single free-parameter rename (ax vs cx).""" + if not a or not b or len(a) != len(b) or len(a) > 80: + return False + diffs = [(i, a[i], b[i]) for i in range(len(a)) if a[i] != b[i]] + if not diffs: + return False + src_chars = set(d[1] for d in diffs) + dst_chars = set(d[2] for d in diffs) + if len(src_chars) == 1 and len(dst_chars) == 1: + src, dst = src_chars.pop(), dst_chars.pop() + if src.isalpha() and dst.isalpha(): + return a.replace(src, dst) == b + return False + + @classmethod + def _normalize_tuple(cls, s: str) -> str: + """Normalize tuple formatting: (2, 5, 609) → 2,5,609.""" + # Strip set-builder conditions: \mid ... or | ... + s = re.sub(r'\\mid.*$', '', s) + s = re.sub(r'\|[^,]*$', '', s) + return re.sub(r'[\s()\[\]{}\\]', '', s) + + @classmethod + def _try_sympy_equal(cls, a: str, b: str) -> bool: + """Optional sympy-based algebraic equivalence (graceful fallback if unavailable).""" + try: + from sympy.parsing.latex import parse_latex + from sympy import simplify, nsimplify + expr_a = parse_latex(a) + expr_b = parse_latex(b) + diff = simplify(nsimplify(expr_a - expr_b)) + return diff == 0 + except Exception: + return False + + @classmethod + def answers_match(cls, predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + + norm_p = cls.normalize_answer(predicted) + norm_r = cls.normalize_answer(reference) + + # --- Strategy 1: direct string equality --- + if norm_p == norm_r: + return True + + # --- Strategy 2: case-insensitive --- + if norm_p.lower() == norm_r.lower(): + return True + + # --- Strategy 3: numeric equality --- + if cls._try_numeric_equal(norm_p, norm_r): + return True + + # --- Strategy 4: variable-prefix stripping (both sides) --- + stripped_p = cls.normalize_answer(cls._strip_var_prefix(predicted)) + stripped_r = cls.normalize_answer(cls._strip_var_prefix(reference)) + if stripped_p and stripped_r and stripped_p == stripped_r: + return True + if stripped_p and stripped_r and cls._try_numeric_equal(stripped_p, stripped_r): + return True + + # --- Strategy 5: MCQ double matching --- + # Extract letter+value from reference + ref_letter, ref_value = cls._extract_mcq_parts(reference) + if ref_letter: + # pred matches the letter? + if norm_p == ref_letter or predicted.strip().upper() == ref_letter: + return True + # pred matches the value? + if ref_value: + norm_ref_val = cls.normalize_answer(ref_value) + if norm_p == norm_ref_val or cls._try_numeric_equal(norm_p, norm_ref_val): + return True + # Extract from predicted side too (pred="B", ref has value) + pred_letter, pred_value = cls._extract_mcq_parts(predicted) + if pred_letter: + if norm_r == pred_letter or reference.strip().upper() == pred_letter: + return True + if pred_value: + norm_pred_val = cls.normalize_answer(pred_value) + if norm_r == norm_pred_val or cls._try_numeric_equal(norm_r, norm_pred_val): + return True + # MCQ: GT is single letter, pred is numeric/expression → match if pred chose option + if cls._MCQ_SINGLE_LETTER_RE.match(reference.strip()): + if cls._MCQ_SINGLE_LETTER_RE.match(predicted.strip().upper()): + return predicted.strip().upper() == reference.strip().upper() + # pred is a value, GT is just a letter: we accept pred=letter match only + # (can't verify value without options text) + + # --- Strategy 6: tuple/set normalization --- + tuple_p = cls._normalize_tuple(norm_p) + tuple_r = cls._normalize_tuple(norm_r) + if ',' in tuple_p and tuple_p == tuple_r: + return True + if stripped_p and stripped_r: + tuple_sp = cls._normalize_tuple(stripped_p) + tuple_sr = cls._normalize_tuple(stripped_r) + if ',' in tuple_sp and tuple_sp == tuple_sr: + return True + + # --- Strategy 7: equation reorder (a+b=c vs c=a+b, or lhs=rhs swapped) --- + if '=' in norm_r and '=' not in norm_p: + # GT has derivation like "18*1+999*2=2016", pred is "2016" + parts = norm_r.split('=') + for part in parts: + part = part.strip() + if part == norm_p or cls._try_numeric_equal(part, norm_p): + return True + if '=' in norm_p and '=' not in norm_r: + parts = norm_p.split('=') + for part in parts: + part = part.strip() + if part == norm_r or cls._try_numeric_equal(part, norm_r): + return True + if '=' in norm_p and '=' in norm_r: + # Both have =: try matching LHS=RHS in any order + pp = [x.strip() for x in norm_p.split('=')] + rp = [x.strip() for x in norm_r.split('=')] + if set(pp) == set(rp): + return True + + # --- Strategy 8: multiplicative reorder (27\pi\sqrt{6} vs 27\sqrt{6}\pi) --- + def _sort_factors(s): + tokens = re.findall(r'\\?[a-zA-Z]+\{[^}]*\}|\\?[a-zA-Z]+|\d+|[^a-zA-Z\d\\{}]', s) + return ''.join(sorted(tokens)) + if _sort_factors(norm_p) == _sort_factors(norm_r): + return True + + # --- Strategy 9: sympy algebraic equivalence (optional, slow) --- + if cls._try_sympy_equal(predicted, reference): + return True + + # --- Strategy 9b: quantifier stripping + param rename --- + q_stripped_r = cls._strip_quantifiers(reference) + q_stripped_p = cls._strip_quantifiers(predicted) + if q_stripped_r != reference or q_stripped_p != predicted: + norm_qr = cls.normalize_answer(cls._strip_var_prefix(q_stripped_r)) + norm_qp = cls.normalize_answer(cls._strip_var_prefix(q_stripped_p)) + if norm_qr and norm_qp: + if norm_qr == norm_qp: + return True + if cls._try_param_rename(norm_qr, norm_qp): + return True + + # --- Strategy 10: param rename on var-prefix-stripped forms --- + if stripped_p and stripped_r and cls._try_param_rename(stripped_p, stripped_r): + return True + + return False + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + user_data = traj.get('user_data') or [] + gt = '' + for item in user_data: + if item[0] == 'ground_truth': + gt = item[1] + break + predicted = self.extract_boxed(completion) + correct = self.answers_match(predicted, gt) + rewards.append(1.0 if correct else 0.0) + return rewards + + +class FormatReward(Reward): + """Reward for having \\boxed{} and ... analysis in the output.""" + + _HINT_RE = re.compile(r'(.*?)', re.DOTALL) + _THINK_RE = re.compile(r'^.*?', re.DOTALL) + _MIN_HINT_LEN = 30 # minimum chars for a substantive hint + _MAX_HINT_LEN = 4096 # hints longer than this are likely thinking dumps + + @staticmethod + def _to_text(content) -> str: + """Convert content (str or list-of-blocks) to plain text.""" + if isinstance(content, str): + return content + if isinstance(content, list): + return ''.join( + b.get('text', '') if isinstance(b, dict) else str(b) + for b in content) + return str(content) if content else '' + + @classmethod + def _visible_response(cls, text: str) -> str: + """Strip ... block to get the visible response.""" + think_end = text.find('') + if think_end >= 0: + return text[think_end + len(''):] + return text + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + sys_content = '' + for msg in messages: + if msg.get('role') == 'system': + sys_content = self._to_text(msg.get('content', '')) + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = self._to_text(msg.get('content', '')) + break + has_boxed = '\\boxed{' in completion + # Only check hint tags for RAG prompts (system contains examples) + is_rag = 'condensed reasoning examples from similar problems' in sys_content + if is_rag: + # Check hint in VISIBLE response only (after ) + visible = self._visible_response(completion) + hint_match = self._HINT_RE.search(visible) + has_good_hint = False + if hint_match: + hint_text = hint_match.group(1).strip() + hint_pos = hint_match.start() + # Hint must be near the start of visible output + at_beginning = hint_pos < max(len(visible) * 0.05, 200) + is_substantive = len(hint_text) >= self._MIN_HINT_LEN + # Reject hints that are too long (model dumping thinking) + not_dump = len(hint_text) <= self._MAX_HINT_LEN + # Must start with the required prefix + has_prefix = hint_text.startswith(HINT_REQUIRED_PREFIX) + has_good_hint = (at_beginning and is_substantive + and not_dump and has_prefix) + # 0.3 for boxed + 0.2 for good hint = 0.5 max + reward = (0.3 if has_boxed else 0.0) + (0.2 if has_good_hint else 0.0) + else: + reward = 0.5 if has_boxed else 0.0 + rewards.append(reward) + return rewards + + +class GibberishPenalty(Reward): + """Negative reward for degenerate outputs (gibberish/random unicode tail).""" + + TAIL_CHARS = 400 + GIBBERISH_THRESHOLD = 0.20 # >20% non-math non-ascii in tail + + @classmethod + def is_gibberish(cls, text: str) -> bool: + if not text: + return False + tail = text[-cls.TAIL_CHARS:] if len(text) > cls.TAIL_CHARS else text + non_math_non_ascii = 0 + for c in tail: + code = ord(c) + # Allow: ASCII, common CJK (for Chinese math), LaTeX symbols + if code > 127 and not (0x4e00 <= code <= 0x9fff): + non_math_non_ascii += 1 + return non_math_non_ascii > len(tail) * cls.GIBBERISH_THRESHOLD + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + rewards.append(-0.5 if self.is_gibberish(completion) else 0.0) + return rewards + + +def compute_rewards(trajectories: List[Dict[str, Any]] + ) -> Tuple[List[float], List[float], List[float]]: + acc_fn = AoPSAccuracyReward() + fmt_fn = FormatReward() + gib_fn = GibberishPenalty() + acc = acc_fn(trajectories) + fmt = fmt_fn(trajectories) + gib = gib_fn(trajectories) + total = [a + f + g for a, f, g in zip(acc, fmt, gib)] + return total, fmt, acc + + +# ============================================================================ +# Dataset: AoPS boxed problems +# ============================================================================ +def create_aops_dataset(): + """Load AoPS and create GRPO-style dataset (prompt only, with ground_truth in user_data).""" + from modelscope import MsDataset + from twinkle.data_format import Message, Trajectory + + ds = MsDataset.load(AOPS_DATASET_ID, split='train', + download_mode='reuse_dataset_if_exists') + rows = [] + for row in ds: + if not row['metadata'].get('boxed'): + continue + ref = AoPSAccuracyReward.extract_boxed(row['solution']) + if not ref: + continue + rows.append({'problem': row['problem'], 'ground_truth': ref}) + + logger.info(f'[aops] loaded {len(rows)} boxed problems') + rng = random.Random(AOPS_SEED) + rng.shuffle(rows) + + # Build Trajectory list (prompt-only for GRPO) + trajectories = [] + for r in rows: + # Use direct system prompt as placeholder — will be replaced by RAG pipeline + traj = Trajectory( + messages=[ + Message(role='system', content=SYSTEM_DIRECT), + Message(role='user', content=r['problem']), + ], + user_data=[('ground_truth', r['ground_truth'])], + ) + trajectories.append(traj) + + data_meta = DatasetMeta(data=trajectories) + dataset = Dataset(data_meta) + dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, + max_length=16384, truncation_strategy='delete', + enable_thinking=True) + dataset.encode(add_generation_prompt=True) + return dataset + + +# ============================================================================ +# Main +# ============================================================================ +def main(): + # GPU rank allocation + cond_start = 0 + emb_start = cond_start + CONDENSER_GPUS + sampler_start = emb_start + EMB_GPUS + model_start = sampler_start + SAMPLER_GPUS + + device_groups = [ + DeviceGroup(name='condenser', ranks=list(range(cond_start, emb_start)), + device_type='GPU'), + DeviceGroup(name='emb_model', ranks=list(range(emb_start, sampler_start)), + device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(sampler_start, model_start)), + device_type='GPU'), + DeviceGroup(name='model', ranks=list(range(model_start, NUM_GPUS)), + device_type='GPU'), + ] + + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=2) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) + condenser_mesh = DeviceMesh.from_sizes(world_size=CONDENSER_GPUS, dp_size=CONDENSER_GPUS) + + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, + groups=device_groups, lazy_collect=False) + + # -- Training model (full-parameter) -- + model = TransformersModel( + model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') + model.set_optimizer('AdamW', lr=LEARNING_RATE) + model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + model.set_loss('GSPOLoss', epsilon=0.2, epsilon_high=0.28, beta=0.04) + model.set_processor(InputProcessor) + model.set_template('Qwen3_5Template', model_id=MODEL_ID, + enable_thinking=True, max_length=32768) + + # -- Rollout sampler -- + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 32768, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, + enable_thinking=True, max_length=32768) + + # -- Embedding model -- + emb_model = TransformersModel( + model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, remote_group='emb_model') + emb_model.set_processor(InputProcessor) + emb_model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) + emb_template = Qwen3_5Template( + model_id=EMBED_MODEL_ID, max_length=EMBED_MAX_LENGTH, + truncation_strategy='delete', enable_thinking=False) + + # -- Condenser sampler -- + condenser_sampler = vLLMSampler( + model_id=CONDENSE_MODEL_ID, + engine_args={'gpu_memory_utilization': 0.85, 'max_model_len': 32768}, + device_mesh=condenser_mesh, + remote_group='condenser', + ) + condenser_sampler.set_template( + 'Qwen3_5Template', model_id=CONDENSE_MODEL_ID, + enable_thinking=False, truncation_strategy='delete', max_length=32768) + condenser_template = Qwen3_5Template( + model_id=CONDENSE_MODEL_ID, max_length=32768, + enable_thinking=False, truncation_strategy='delete') + condenser_special_tokens = set(condenser_template.tokenizer.all_special_tokens) + compress_params = SamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, temperature=CONDENSE_TEMPERATURE, + top_p=0.5, num_samples=1) + + # -- API client (condenser fallback) -- + api_client = None + if CONDENSE_API_KEY: + api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + + # -- LanceDB -- + import lancedb + db = lancedb.connect(DB_PATH) + tbl = db.open_table(DB_TABLE) + logger.info(f'[rag] LanceDB ready, rows={tbl.count_rows()}') + + # -- Checkpoint & DataLoader -- + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + + GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS + dataloader = DataLoader( + dataset=create_aops_dataset, + batch_size=GLOBAL_BATCH_SIZE, + min_batch_size=GLOBAL_BATCH_SIZE, + device_mesh=model_mesh, + remote_group='model', + ) + + advantage_fn = GRPOAdvantage() + metrics = CompletionRewardMetric() + sampling_params = SamplingParams( + max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, + temperature=1.0, top_p=0.95) + + optim_step = 0 + logger.info('Starting RAG-hint GRPO training') + logger.info(get_device_placement()) + + # -- Prefetch: overlap RAG data preparation with training -- + prefetch_pool = ThreadPoolExecutor(max_workers=1) + + def _extract_text(content) -> str: + """Extract plain text from content (str or list-of-parts format).""" + if isinstance(content, str): + return content + if isinstance(content, list): + return ''.join( + p.get('text', '') for p in content if isinstance(p, dict) and p.get('type') == 'text') + return str(content) if content else '' + + def prepare_rag_batch(batch): + """Embed → retrieve → condense → build prompts. Runs in background thread.""" + problems = [] + ground_truths = [] + for item in batch: + msgs = item.get('messages', []) + prob = '' + for m in msgs: + if m.get('role') == 'user': + prob = _extract_text(m.get('content', '')) + break + problems.append(prob) + ud = item.get('user_data', []) + gt = '' + for pair in ud: + if pair[0] == 'ground_truth': + gt = pair[1] + break + ground_truths.append(gt) + + # Embed & retrieve + query_vecs = get_embeddings(emb_model, emb_template, problems, EMB_GPUS) + retrieved = retrieve_topk(tbl, query_vecs, problems, SIM_THRESHOLD) + raw_retrieved_counts = [len(r) for r in retrieved] + + # LLM-based decontamination: judge ALL retrievals via API + if api_client: + judge_pairs = [] # (qi, ret_idx, prob_a, prob_b) + for qi, rets in enumerate(retrieved): + for ri, ret in enumerate(rets): + judge_pairs.append((qi, ri, problems[qi], ret['query'])) + + if judge_pairs: + pairs_input = [(pa, pb) for _, _, pa, pb in judge_pairs] + verdicts = _llm_judge_same_problem(api_client, pairs_input) + to_remove = set() + for vi, (qi, ri, _, _) in enumerate(judge_pairs): + if verdicts[vi]: + to_remove.add((qi, ri)) + if to_remove: + logger.info(f'[decontam-llm] filtered {len(to_remove)} same-problem retrievals') + for qi in range(len(retrieved)): + retrieved[qi] = [ + ret for ri, ret in enumerate(retrieved[qi]) + if (qi, ri) not in to_remove + ] + + # Condense (batch local vLLM + API fallback) + condensed_examples: List[List[Dict[str, str]]] = [[] for _ in range(len(problems))] + tasks_to_condense = [] + for i, rets in enumerate(retrieved): + for j, ret in enumerate(rets): + tasks_to_condense.append((i, j, problems[i], ret)) + + if tasks_to_condense: + condense_prompts = [] + for idx, _j, prob, ret in tasks_to_condense: + user_msg = COMPRESS_USER.format(query=prob, text=ret['thinking']) + condense_prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_msg}]}) + + try: + condense_responses = condenser_sampler.sample(condense_prompts, compress_params) + except Exception as exc: + logger.warning(f'[condense] local batch error: {exc}') + condense_responses = [None] * len(condense_prompts) + + api_fallback_indices = [] + for ci, (idx, _j, prob, ret) in enumerate(tasks_to_condense): + resp = condense_responses[ci] if condense_responses else None + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + for tok in condenser_special_tokens: + text = text.replace(tok, '') + text = text.strip() + if text: + condensed_examples[idx].append({'query': ret['query'], 'thinking': text}) + else: + api_fallback_indices.append(ci) + + if api_fallback_indices and api_client: + def _fallback(ci): + return ci, _api_condense_single(api_client, condense_prompts[ci]['messages']) + with ThreadPoolExecutor(max_workers=CONDENSE_API_CONCURRENCY) as pool: + futs = [pool.submit(_fallback, ci) for ci in api_fallback_indices] + for fut in as_completed(futs): + ci, result = fut.result() + idx, _j, prob, ret = tasks_to_condense[ci] + text = result if result else ret['thinking'][:MAX_TRACE_LEN] + condensed_examples[idx].append({'query': ret['query'], 'thinking': text}) + elif api_fallback_indices: + for ci in api_fallback_indices: + idx, _j, prob, ret = tasks_to_condense[ci] + condensed_examples[idx].append( + {'query': ret['query'], 'thinking': ret['thinking'][:MAX_TRACE_LEN]}) + + # API hint analysis: pre-compute RAG relevance verdict + hint_analyses = [None] * len(problems) + if api_client: + hint_analyses = _api_hint_analysis_batch(api_client, problems, condensed_examples) + + # Build prompts with rag_fallback_sim check + rag_prompts = [] + rag_debug_records = [] + for i, prob in enumerate(problems): + examples = condensed_examples[i] + rets = retrieved[i] + best_sim = max((r['sim'] for r in rets), default=0.0) + use_rag = bool(examples) and best_sim >= RAG_FALLBACK_SIM + + if use_rag: + # If API hint analysis succeeded, use pre-analyzed prompt (no needed) + if hint_analyses[i]: + rag_sys_content = build_preanalysis_system(hint_analyses[i]) + else: + # Fallback: old-style prompt with self-analysis requirement + parts = [SYSTEM_WITH_RAG_HEADER] + for eidx, ex in enumerate(examples, 1): + parts.append(EXAMPLE_TEMPLATE.format( + idx=eidx, + example_query=ex['query'], + example_thinking=ex['thinking'])) + rag_sys_content = ''.join(parts) + + # RAG group (only RAG, no paired NoRAG) + rag_prompts.append({ + 'messages': [ + {'role': 'system', 'content': rag_sys_content}, + {'role': 'user', 'content': prob}, + ], + 'user_data': [('ground_truth', ground_truths[i])], + 'assistant_prefix': ANALYSIS_PREFIX, + }) + rag_debug_records.append({ + 'problem': prob[:200], + 'ground_truth': ground_truths[i], + 'best_sim': round(best_sim, 4), + 'num_raw_retrieved': raw_retrieved_counts[i], + 'num_retrieved': len(rets), + 'num_condensed': len(examples), + 'use_rag': True, + 'has_preanalysis': hint_analyses[i] is not None, + 'preanalysis_len': len(hint_analyses[i]) if hint_analyses[i] else 0, + 'top_retrieved_query': rets[0]['query'][:200] if rets else '', + 'condensed_len': len(examples[0].get('thinking', '')) if examples else 0, + }) + else: + # No hint found — skip this query entirely + continue + + return rag_prompts, rag_debug_records + + # Submit first batch prefetch + os.makedirs(OUTPUT_DIR, exist_ok=True) + rag_log_path = os.path.join(OUTPUT_DIR, 'rag_diagnostics.jsonl') + rag_log_f = open(rag_log_path, 'w', encoding='utf-8') + logger.info(f'[rag] diagnostics → {rag_log_path}') + + batch_iter = iter(dataloader) + pending_future = None + try: + first_batch = next(batch_iter) + pending_future = prefetch_pool.submit(prepare_rag_batch, first_batch) + except StopIteration: + pass + + try: + while pending_future is not None: + if optim_step >= MAX_STEPS: + break + + metrics.reset() + rag_prompts, rag_debug_records = pending_future.result() + + # Write RAG diagnostics + for rec in rag_debug_records: + rec['step'] = optim_step + rag_log_f.write(json.dumps(rec, ensure_ascii=False) + '\n') + rag_log_f.flush() + + # Submit next batch prefetch (overlaps with rollout + training) + pending_future = None + try: + next_batch = next(batch_iter) + pending_future = prefetch_pool.submit(prepare_rag_batch, next_batch) + except StopIteration: + pass + + # ---- Expand for NUM_GENERATIONS and sample ---- + expand_prompts = [] + for prompt in rag_prompts: + expand_prompts.extend([prompt] * NUM_GENERATIONS) + + if not expand_prompts: + logger.warning(f'[Step {optim_step}] empty prompt list after RAG processing, skip') + continue + + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + + sample_responses = sampler.sample(expand_prompts, sampling_params) + + # ---- Collect rollouts ---- + all_input_data: List[Dict[str, Any]] = [] + all_old_logps: List[List[float]] = [] + all_completion_lengths: List[int] = [] + + for sample_response in sample_responses: + for sequence in sample_response.sequences: + all_input_data.append(sequence.new_input_feature) + all_old_logps.append([logprob[0][1] for logprob in sequence.logprobs]) + all_completion_lengths.append(len(sequence.tokens)) + + # ---- Rewards ---- + total_rewards, format_rewards, accuracy_rewards = compute_rewards(all_input_data) + + # Zero out rewards for rollouts that hit the max_tokens ceiling + max_len_threshold = int(MAX_NEW_TOKENS * 0.95) + for i in range(len(all_input_data)): + if all_completion_lengths[i] >= max_len_threshold: + total_rewards[i] = 0.0 + accuracy_rewards[i] = 0.0 + format_rewards[i] = 0.0 + + # Per-step reward summary to diagnostics + n_correct = sum(1 for a in accuracy_rewards if a > 0) + rag_log_f.write(json.dumps({ + 'step': optim_step, 'type': 'reward_summary', + 'n_samples': len(accuracy_rewards), + 'accuracy': n_correct / len(accuracy_rewards) if accuracy_rewards else 0, + 'mean_reward': sum(total_rewards) / len(total_rewards) if total_rewards else 0, + }, ensure_ascii=False) + '\n') + + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={ + 'total': total_rewards, + 'format': format_rewards, + 'accuracy': accuracy_rewards, + }, + ) + + # ---- GRPO advantage ---- + advantages = advantage_fn( + total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + if ADV_CLIP > 0: + advantages = [max(-ADV_CLIP, min(ADV_CLIP, a)) for a in advantages] + + # Log all rollout responses (after advantage computation) + _extract_boxed = AoPSAccuracyReward.extract_boxed + def _content_to_str(content): + """Convert message content (str or list of blocks) to plain text.""" + if isinstance(content, str): + return content + if isinstance(content, list): + return ''.join( + b.get('text', '') if isinstance(b, dict) else str(b) + for b in content) + return str(content) + + for ridx, traj in enumerate(all_input_data): + msgs = traj.get('messages', []) + assistant_text = _content_to_str(next( + (m['content'] for m in reversed(msgs) if m.get('role') == 'assistant'), '')) + user_text = _content_to_str(next( + (m['content'] for m in msgs if m.get('role') == 'user'), '')) + sys_text = _content_to_str(next( + (m['content'] for m in msgs if m.get('role') == 'system'), '')) + user_data = traj.get('user_data') or [] + gt = next((v for k, v in user_data if k == 'ground_truth'), '') + problem_idx = ridx // NUM_GENERATIONS + use_rag = ('condensed reasoning examples from similar problems' in sys_text + or 'RAG Analysis (pre-computed)' in sys_text) + # Per-problem group accuracy (all generations for same problem) + grp_start = problem_idx * NUM_GENERATIONS + grp_end = grp_start + NUM_GENERATIONS + grp_acc = sum(accuracy_rewards[grp_start:grp_end]) / NUM_GENERATIONS + + rag_log_f.write(json.dumps({ + 'step': optim_step, 'type': 'rollout', + 'idx': ridx, + 'problem_idx': problem_idx, + 'problem': user_text, + 'system': sys_text, + 'response': assistant_text, + 'ground_truth': gt, + 'predicted': _extract_boxed(assistant_text), + 'use_rag': use_rag, + 'best_sim': rag_debug_records[problem_idx].get('best_sim', 0.0) if problem_idx < len(rag_debug_records) else 0.0, + 'reward': total_rewards[ridx], + 'accuracy_reward': accuracy_rewards[ridx], + 'format_reward': format_rewards[ridx], + 'advantage': advantages[ridx], + 'completion_length': all_completion_lengths[ridx], + 'group_accuracy': grp_acc, + }, ensure_ascii=False) + '\n') + + rag_log_f.flush() + + # ---- Filter out low-signal problem groups (DAPO-style dynamic sampling) ---- + # Skip groups where accuracy is too low (<0.1) or too high (>0.9) + # to avoid gradient dominated by gibberish/format noise or no learning signal. + filtered_inputs, filtered_old_logps, filtered_advantages = [], [], [] + actual_num_groups = len(all_input_data) // NUM_GENERATIONS + for g in range(actual_num_groups): + g_start = g * NUM_GENERATIONS + g_end = g_start + NUM_GENERATIONS + grp_adv = advantages[g_start:g_end] + if all(abs(a) < 1e-8 for a in grp_adv): + continue + grp_acc_rate = sum(accuracy_rewards[g_start:g_end]) / NUM_GENERATIONS + if grp_acc_rate < 0.2 or grp_acc_rate > 0.8: + continue + filtered_inputs.extend(all_input_data[g_start:g_end]) + filtered_old_logps.extend(all_old_logps[g_start:g_end]) + filtered_advantages.extend(grp_adv) + + # ---- Mini-batch training with gradient accumulation ---- + # Process MICRO_BATCH_SIZE samples per forward, accumulate grad_accum_steps + # times before one optimizer step. clip_grad_norm normalizes by accumulated + # num_tokens, ensuring mathematical equivalence with larger batch forward. + total_completions = len(filtered_inputs) + if total_completions == 0: + logger.info(f'[Step {optim_step}] all groups filtered (uniform rewards), skip training') + continue + + grad_accum_steps = MINI_BATCH_SIZE // MICRO_BATCH_SIZE + accum_count = 0 + for mb_start in range(0, total_completions, MICRO_BATCH_SIZE): + mb_end = min(mb_start + MICRO_BATCH_SIZE, total_completions) + mb_inputs = filtered_inputs[mb_start:mb_end] + mb_old_logps = filtered_old_logps[mb_start:mb_end] + mb_advantages = filtered_advantages[mb_start:mb_end] + + outputs = model.forward_backward( + inputs=mb_inputs, + old_logps=mb_old_logps, + ref_logps=mb_old_logps, + advantages=mb_advantages, + ) + accum_count += 1 + + if accum_count % grad_accum_steps == 0: + # Loss spike skip: discard explosive gradients + skip_step = False + try: + loss_val = outputs.get('loss', None) + if loss_val is not None: + if hasattr(loss_val, 'item'): + loss_val = loss_val.item() + if loss_val > LOSS_SPIKE_THRESHOLD: + skip_step = True + logger.warning( + f'[Step {optim_step}] Loss spike: {loss_val:.4f} > ' + f'{LOSS_SPIKE_THRESHOLD}, skipping update') + except Exception: + pass + + if skip_step: + model.zero_grad() + else: + model.clip_grad_and_step() + optim_step += 1 + + if optim_step >= MAX_STEPS: + break + if optim_step % SAVE_STEPS == 0: + model.save(f'rag-hint-grpo-checkpoint-{optim_step}') + + # Flush remaining accumulated gradients (incomplete window at tail) + if accum_count % grad_accum_steps != 0: + skip_step = False + try: + loss_val = outputs.get('loss', None) + if loss_val is not None: + if hasattr(loss_val, 'item'): + loss_val = loss_val.item() + if loss_val > LOSS_SPIKE_THRESHOLD: + skip_step = True + logger.warning( + f'[Step {optim_step}] Loss spike (tail): {loss_val:.4f} > ' + f'{LOSS_SPIKE_THRESHOLD}, skipping update') + except Exception: + pass + + if skip_step: + model.zero_grad() + else: + model.clip_grad_and_step() + optim_step += 1 + + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True)) + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + finally: + prefetch_pool.shutdown(wait=False) + rag_log_f.close() + + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('rag-hint-grpo-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/sample/emb_sample.py b/cookbook/sample/emb_sample.py index da27a815..8db4b91a 100644 --- a/cookbook/sample/emb_sample.py +++ b/cookbook/sample/emb_sample.py @@ -32,10 +32,10 @@ args = CLI.from_args() # -- Config ------------------------------------------------------------------- -CONDENSE_MODEL_ID = args.extra.get('condense_model_id', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') -EMB_MODEL_ID = args.extra.get('emb_model_id', 'ms://twinkle-kit/Qwen3.5-4B-QA-emb') -SAMPLER_GPUS = args.infra.sampler_gpus or 1 -EMB_GPUS = int(args.extra.get('emb_gpus', 1)) +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +EMB_MODEL_ID = os.environ.get('EMB_MODEL', 'output/embedding_lora_transformers/step_8000') +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 1)) +EMB_GPUS = int(os.environ.get('EMB_GPUS', 1)) EMB_MAX_LENGTH = 8192 # -- Prompts (aligned with train_embedding_full_ddp.py) ----------------------- diff --git a/cookbook/sample/rag_recall_sample.py b/cookbook/sample/rag_recall_sample.py new file mode 100644 index 00000000..691a69f6 --- /dev/null +++ b/cookbook/sample/rag_recall_sample.py @@ -0,0 +1,379 @@ +"""RAG recall test: compress a query via condenser → embed → search LanceDB. + +End-to-end validation that the thinking-trace RAG index built by +``cookbook/exp/embedding/build_thinking_rag_index.py`` is retrievable. + +Architecture (8 GPUs, same as build script): + * GPU 0-3: vLLM condenser (TP=4) + * GPU 4-7: TransformersModel embedding (DP=4) + +Launch: + python cookbook/sample/rag_recall_sample.py + python cookbook/sample/rag_recall_sample.py --query "How to implement binary search?" + python cookbook/sample/rag_recall_sample.py --db-path ./output/thinking_rag/lance.db --top-k 5 +""" +import argparse +import os +import re +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.loss import InfonceLoss +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template + +logger = get_logger() + +# --------------------------------------------------------------------------- +# Config (mirrors build_thinking_rag_index.py) +# --------------------------------------------------------------------------- +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +EMBED_MODEL_ID = os.environ.get( + 'EMBED_MODEL_ID', 'output/embedding_lora_transformers/last-checkpoint') +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 4)) +EMB_GPUS = int(os.environ.get('EMB_GPUS', 4)) +NUM_GPUS = SAMPLER_GPUS + EMB_GPUS + +CONDENSE_GPU_MEM = float(os.environ.get('CONDENSE_GPU_MEM', 0.85)) +CONDENSE_MAX_MODEL_LEN = int(os.environ.get('CONDENSE_MAX_MODEL_LEN', 32768)) +CONDENSE_MAX_TOKENS = int(os.environ.get('CONDENSE_MAX_TOKENS', 8192)) +COMPRESS_TEMPERATURE = float(os.environ.get('COMPRESS_TEMPERATURE', 0.2)) +COMPRESS_TOP_P = float(os.environ.get('COMPRESS_TOP_P', 0.5)) +EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 8192)) +MIN_TEXT_CHARS = int(os.environ.get('MIN_TEXT_CHARS', 256)) + +# --------------------------------------------------------------------------- +# Compress prompts — MUST match build_thinking_rag_index.py exactly. +# --------------------------------------------------------------------------- +COMPRESS_SYSTEM = """\ +You are a compression and summary assistant. For the (query, source) pair, emit a Markdown \ +answer with TWO sections, designed to pair with the `extract_compressed` tool: \ +the reader absorbs `## Summary` directly, then calls `extract_compressed` \ +on any topic-key listed under `## More` to recover its \ +fuller content. + + `## Summary` \u2014 extreme-density text the reader reads directly. + `## More` \u2014 a topic index whose keys are valid arguments \ +to `extract_compressed` for recovering material not captured inline. + +Together the two sections must form a COMPLETE, NON-DISTORTING inventory of the \ +source for the query \u2014 nothing essential lost, nothing implied that the source \ +does not support. NO preamble, NO meta-commentary, NO code fences wrapping the \ +whole output. + +Output skeleton: + +## Summary +Topic: + + +## More +- : +- ... + +Format selection for the inline body (pick the MOST COMPACT form per query, mix \ +when helpful): +- Interface / signature \u2192 code notation directly: `func(a:int)->str` +- Factual / entity \u2192 telegraphic prose; drop function words; \":\" for \"is\", \",\" \ +for \"has\" +- Skill / how-to / usage \u2192 lead with `Use when: `; numbered telegraphic \ +steps `1.do X 2.then Y`; close with `Output: ` when relevant +- Procedural \u2192 numbered short steps +- Analytical / design \u2192 hierarchical bullets with abbreviations + +`## Summary` rules: +1. TOPIC LINE \u2014 line 1 is ALWAYS `Topic: `, even when the \ +query is narrow. Anchors both the reader and the tool. +2. DENSITY \u2014 every token in the body carries query-relevant signal; cut filler. +3. PRIMARY-COMPLETE \u2014 never silently drop a fact essential to answering the \ +query. Anything cut for length MUST appear as a key under \ +`## More`. +4. NON-MISLEADING \u2014 phrasing must not let the reader infer anything the source \ +does not support; partial truths that mislead are worse than honest omissions \ +flagged in the index. +5. SELF-CONTAINED \u2014 the reader can act on the answer without re-opening the source. +6. FAITHFUL \u2014 only content the source supports; no fabrication, no extrapolation. +7. LANGUAGE \u2014 match the source language. +8. NO outer code fences around the whole answer; no meta-commentary. + +`## More` rules (MANDATORY \u2014 this section is never omitted): +1. FORMAT \u2014 each bullet is `- : `: + \u2022 topic-key \u2014 short, unambiguous, grounded in source vocabulary so the \ +`extract_compressed` tool can locate the aspect (e.g. `decorators`, \ +`error handling`, `pitfalls`). + \u2022 hint \u2014 tells WHAT the reader gains by expanding (concrete numbers, code \ +listings, secondary cases, edge details, related context, \u2026); do NOT restate \ +the inline answer. +2. CRITERION \u2014 each bullet names an aspect that EXISTS in the source but is \ +NOT fully captured inline. Material that genuinely fits inline without \ +distortion MUST NOT be duplicated here. +3. FAITHFUL \u2014 hints must be grounded in the source; never speculate or invent. +4. ORDER \u2014 by relevance to the query, then by importance. +5. EMPTY CASE \u2014 if the source is so short / single-purpose that everything \ +fits inline, write a single line `- (none)`. + +Now begin.\ +""" + +COMPRESS_USER = ( + 'Downstream model will read your compressed block to decide whether to ' + 'expand it. Compress faithfully: preserve the passage topic + core facts. ' + 'Do NOT invent facts. Do NOT drop major facts. Do NOT write meta-commentary ' + 'about the Query (never write "Query info: absent", "no X mention", etc.); ' + 'if the passage does not address the Query, still summarize the passage. ' + 'CRITICAL LANGUAGE RULE: detect the dominant language of the Passage ' + '(NOT the Query, NOT this instruction) and write the ENTIRE output in that ' + 'same language; English passage \u2192 English output, Chinese passage \u2192 ' + 'Chinese output, Japanese passage \u2192 Japanese output. NEVER translate, ' + 'NEVER mix languages, NEVER copy these instructions into the output.\n\n' + '## Query (ordering hint only \u2014 still summarize the whole passage)\n{query}\n\n' + '## Passage\n{text}') + +RAG_QUERY_HINT = ( + 'Summarize this query for retrieval. ' + 'The body of ## Summary MUST follow this EXACT 4-line template \u2014 ' + 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the specific pattern, never generic labels.') + +# --------------------------------------------------------------------------- +# Demo queries (diverse domains to exercise retrieval) +# --------------------------------------------------------------------------- +DEMO_QUERIES = [ + 'How can I implement binary search in Python and what are the edge cases?', + 'Explain the Free-Energy Principle in neuroscience and how it relates to active inference.', + '如何用动态规划解决最长公共子序列问题？', + 'What is the optimal turbulence model for simulating airflow around a building?', + '请详细解释快速排序的分治策略及其时间复杂度分析', +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _strip_outer_codefence(text: str) -> str: + m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', text, re.DOTALL) + return m.group(1).strip() if m else text.strip() + + +def _short(text: str, n: int = 120) -> str: + text = (text or '').replace('\n', ' ').strip() + return text[:n] + ('\u2026' if len(text) > n else '') + + +def _build_compress_messages(text: str, query: str) -> List[Dict[str, str]]: + return [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': COMPRESS_USER.format(query=query, text=text)}, + ] + + +def _wrap_anchor(text: str) -> List[Dict[str, str]]: + return [ + {'role': 'user', 'content': text}, + {'role': 'assistant', 'content': 'Match the correct response here.'}, + ] + + +# --------------------------------------------------------------------------- +# Core pipeline +# --------------------------------------------------------------------------- + +def compress_query(sampler: vLLMSampler, query: str) -> str: + """Compress a query using the condenser; short queries pass through.""" + if len(query) < MIN_TEXT_CHARS: + return query + prompts = [{'messages': _build_compress_messages(query, RAG_QUERY_HINT)}] + params = SamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, + temperature=COMPRESS_TEMPERATURE, + top_p=COMPRESS_TOP_P, + num_samples=1, + ) + responses = sampler.sample(prompts, params) + seq = responses[0].sequences[0] if responses and responses[0].sequences else None + if seq is None: + return query + text = seq.decoded or '' + text = re.sub(r'<\|[^|]+\|>', '', text).rstrip() + text = _strip_outer_codefence(text) + return text if text.strip() else query + + +def embed_query(model: TransformersModel, template: Qwen3_5Template, + text: str) -> np.ndarray: + """Encode a single text as an anchor embedding, returns [H] float32.""" + feat = template.encode({'messages': _wrap_anchor(text)}) + feat['labels'] = [1] + # Pad to EMB_GPUS to avoid dispatch starvation. + pad_n = EMB_GPUS - 1 + pad_feat = template.encode({'messages': _wrap_anchor(' ')}) + pad_feat['labels'] = [1] + features = [feat] + [pad_feat] * pad_n + out = model.forward_only(inputs=features, task='embedding', return_logits=True) + emb = out['embeddings'] + if hasattr(emb, 'detach'): + emb = emb.detach().cpu().numpy() + return np.asarray(emb[0], dtype=np.float32) + + +def search_lancedb(db_path: str, table_name: str, vector: np.ndarray, + top_k: int) -> List[Dict[str, Any]]: + """Search LanceDB table and return top-k results.""" + import lancedb + db = lancedb.connect(db_path) + available = db.list_tables() + table_list = available.tables if hasattr(available, 'tables') else list(available) + if table_name not in table_list: + raise SystemExit(f'Table "{table_name}" not found in {db_path}. ' + f'Available: {table_list}') + tbl = db.open_table(table_name) + results = ( + tbl.search(vector.tolist()) + .metric('dot') + .limit(top_k) + .select(['id', 'source', 'query_raw', 'thinking_raw', + 'query_compressed', 'cot_compressed', 'sim', '_distance']) + .to_list() + ) + return results + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--query', type=str, nargs='*', default=None, + help='Custom queries to test (overrides built-in demos).') + p.add_argument('--db-path', default='./output/thinking_rag/lance.db', + help='LanceDB directory (same as build script).') + p.add_argument('--table', default='thinking_traces', + help='LanceDB table name.') + p.add_argument('--top-k', type=int, default=3, + help='Number of results to retrieve per query.') + return p.parse_args() + + +def main(): + args = parse_args() + + if not Path(args.db_path).exists(): + raise SystemExit(f'DB path does not exist: {args.db_path}\n' + f'Run build_thinking_rag_index.py first.') + + queries = args.query if args.query else DEMO_QUERIES + + # ── 1. Initialize Twinkle ─────────────────────────────────────────── + device_groups = [ + DeviceGroup( + name='sampler', + ranks=list(range(SAMPLER_GPUS)), + device_type='GPU', + gpus_per_worker=SAMPLER_GPUS, + ), + DeviceGroup( + name='emb_model', + ranks=list(range(SAMPLER_GPUS, NUM_GPUS)), + device_type='GPU', + ), + ] + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, tp_size=SAMPLER_GPUS) + emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) + twinkle.initialize( + mode='ray', nproc_per_node=NUM_GPUS, + groups=device_groups, lazy_collect=False) + + # ── 2. vLLM condenser ─────────────────────────────────────────────── + sampler = vLLMSampler( + model_id=CONDENSE_MODEL_ID, + engine_args={ + 'gpu_memory_utilization': CONDENSE_GPU_MEM, + 'max_model_len': CONDENSE_MAX_MODEL_LEN, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template( + 'Qwen3_5Template', model_id=CONDENSE_MODEL_ID, + enable_thinking=False, max_length=CONDENSE_MAX_MODEL_LEN) + + # ── 3. Embedding model ────────────────────────────────────────────── + emb_model = TransformersModel( + model_id=EMBED_MODEL_ID, + device_mesh=emb_mesh, + remote_group='emb_model', + ) + emb_model.set_processor(InputProcessor) + emb_model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) + emb_template = Qwen3_5Template( + model_id=EMBED_MODEL_ID, + max_length=EMBED_MAX_LENGTH, + truncation_strategy='delete', + enable_thinking=False, + ) + + logger.info(f'Initialized: sampler GPUs 0-{SAMPLER_GPUS-1}, ' + f'emb GPUs {SAMPLER_GPUS}-{NUM_GPUS-1}') + logger.info(f'DB: {args.db_path} / table: {args.table}') + logger.info(f'Queries to test: {len(queries)}') + + # ── 4. Per-query: compress → embed → search ───────────────────────── + for i, raw_query in enumerate(queries): + print(f'\n{"="*80}') + print(f'[Query {i+1}/{len(queries)}]') + print(f' Raw: {_short(raw_query, 200)}') + + # Compress + compressed = compress_query(sampler, raw_query) + is_passthrough = len(raw_query) < MIN_TEXT_CHARS + if is_passthrough: + print(f' Compressed: (passthrough, len={len(raw_query)} < {MIN_TEXT_CHARS})') + else: + print(f' Compressed ({len(raw_query)}\u2192{len(compressed)} chars):') + for line in compressed.split('\n')[:8]: + print(f' {line}') + if compressed.count('\n') > 8: + print(f' ... ({compressed.count(chr(10))+1} lines total)') + + # Embed + vec = embed_query(emb_model, emb_template, compressed) + print(f' Embedding: shape={vec.shape}, norm={np.linalg.norm(vec):.4f}') + + # Search + results = search_lancedb(args.db_path, args.table, vec, args.top_k) + print(f'\n Top-{args.top_k} Results:') + if not results: + print(' (no results)') + continue + for rank, r in enumerate(results, 1): + dist = r.get('_distance', None) + sim = (1.0 - dist) if isinstance(dist, (int, float)) else None + sim_str = f'{sim:.4f}' if sim is not None else '?' + dist_str = f'{dist:.4f}' if isinstance(dist, (int, float)) else '?' + print(f' [{rank}] cos_sim={sim_str} (dist={dist_str}) source={r["source"]}') + print(f' query: {_short(r["query_raw"], 100)}') + print(f' thinking: {_short(r["thinking_raw"], 150)}') + print() + + print(f'\n{"="*80}') + print('RAG recall test complete.') + + +if __name__ == '__main__': + main() diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 781b2206..7fb799ec 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -64,8 +64,8 @@ def _compute_log_importance_weights( """ import torch log_ratio = per_token_logps - per_token_old_logps - # Clamp for numerical stability - log_ratio = torch.clamp(log_ratio, min=-20.0, max=20.0) + # Clamp for numerical stability (±5 bounds ratio to [exp(-5), exp(5)] ≈ [0.007, 148]) + log_ratio = torch.clamp(log_ratio, min=-5.0, max=5.0) return log_ratio def _compute_per_token_loss( @@ -75,9 +75,16 @@ def _compute_per_token_loss( per_token_logps: 'torch.Tensor', ) -> 'torch.Tensor': """ - Compute per-token loss with PPO clipping. + Compute per-token loss with PPO double-sided clipping. - Override this method in subclasses for different loss formulations. + Standard PPO clip is one-sided: it only bounds the loss when + advantage > 0 and ratio > 1+eps. When advantage < 0 and ratio > 1+eps + (policy moved AWAY from the old action on its own), the loss is + unbounded upward, causing gradient explosions. + + This implementation clips the ratio from BOTH sides regardless of + advantage sign, bounding the per-token loss to at most + (1+eps_high) * |advantage|. Args: ratio: [batch, seq_len] importance sampling ratio @@ -91,7 +98,14 @@ def _compute_per_token_loss( clipped_ratio = torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon_high) loss1 = ratio * advantages loss2 = clipped_ratio * advantages - return -torch.min(loss1, loss2) + # Double-sided clip: use max for positive advantage, min for negative. + # Equivalent to: always take the MORE conservative (smaller magnitude) loss. + per_token_loss = torch.where( + advantages >= 0, + -torch.min(loss1, loss2), # positive adv: standard PPO clip + -torch.max(loss1, loss2), # negative adv: clip the OTHER side + ) + return per_token_loss def _aggregate_loss( self, @@ -320,7 +334,7 @@ def _compute_log_importance_weights( """Sequence-level importance sampling: use mean log ratio.""" import torch log_ratio = per_token_logps - per_token_old_logps - log_ratio = torch.clamp(log_ratio, min=-20.0, max=20.0) + log_ratio = torch.clamp(log_ratio, min=-5.0, max=5.0) seq_level_log_weights = ((log_ratio * loss_mask).sum(-1) / loss_mask.sum(-1).clamp(min=1.0)).unsqueeze(-1) return seq_level_log_weights diff --git a/src/twinkle/template/base.py b/src/twinkle/template/base.py index 3c6c29f6..ae641111 100644 --- a/src/twinkle/template/base.py +++ b/src/twinkle/template/base.py @@ -682,7 +682,10 @@ def encode(self, trajectory: Trajectory, add_generation_prompt: bool = False, ** assert self.truncation_strategy != 'split', ( 'encode() does not support truncation_strategy=="split" because it may produce multiple outputs. ' 'Use batch_encode() instead.') - return self.batch_encode([trajectory], add_generation_prompt=add_generation_prompt, **kwargs)[0] + encoded = self.batch_encode([trajectory], add_generation_prompt=add_generation_prompt, **kwargs) + if encoded: + return encoded[0] + return None @staticmethod def map_col_to_row(trajectories: Dict[str, Any]):