From 05f84d93a63617a882f8c4245968357887b4e5e9 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Sun, 14 Jun 2026 22:39:32 +0800 Subject: [PATCH 01/26] fix --- .../exp/embedding/build_thinking_rag_index.py | 810 ++++++++++++++++++ cookbook/exp/embedding/dataset_index.py | 718 ++++++++++++++++ .../exp/embedding/train_embedding_full_ddp.py | 49 +- 3 files changed, 1566 insertions(+), 11 deletions(-) create mode 100644 cookbook/exp/embedding/build_thinking_rag_index.py create mode 100644 cookbook/exp/embedding/dataset_index.py diff --git a/cookbook/exp/embedding/build_thinking_rag_index.py b/cookbook/exp/embedding/build_thinking_rag_index.py new file mode 100644 index 00000000..cf4d1dd7 --- /dev/null +++ b/cookbook/exp/embedding/build_thinking_rag_index.py @@ -0,0 +1,810 @@ +"""Build a thinking-trace RAG index from condensed (query, cot) pairs. + +Pipeline (per row, batched): + 1. Load (user_query, reasoning_content) pairs from ``dataset_think.get_dataset``. + 2. Compress query with ``FIXED_QUERY_NEED`` and cot with ``FIXED_QUERY_SKILL`` + using a Twinkle ``vLLMSampler`` (TP=4 across GPUs 0-3). Reuses the prompt + suite from ``cookbook/exp/condenser/make_condenser_dataset.py``. + 3. On condenser truncation (``stop_reason='length'`` or skeleton-incomplete + output), fall back to an external OpenAI-compatible API. + 4. Encode the condensed pair via the trained embedding model — Twinkle + ``TransformersModel`` on the ``emb_model`` device group (DP=4 across GPUs + 4-7) using ``forward_only(task='embedding')``, the same code path as + training. + 5. Compute cosine similarity for each (query, thinking) pair, drop pairs with + ``sim < SIM_THRESHOLD``, and insert kept rows into LanceDB. The vector + column carries the **positive (compressed-skill)** embedding so a search + keyed by an anchor-encoded query retrieves the matching thinking trace. + 6. Each row stores the **raw thinking** alongside its embedding, so a hit + in the index can directly surface the original CoT. + +Eval mode (``--mode eval`` or ``--mode both``): + * Self-recall test — encode a sample of dataset queries (whose corresponding + rows are already in the index) as anchors and report recall@1/5/10 plus + a per-source breakdown. + +Architecture (8 GPUs): + * GPU 0-3: vLLM condenser (tensor-parallel, ``DeviceGroup name='sampler'``) + * GPU 4-7: TransformersModel embedding (data-parallel, ``DeviceGroup name='emb_model'``) + * Single ``twinkle.initialize(mode='ray', ...)`` call wires both groups. + +Launch examples: + python build_thinking_rag_index.py --mode build --total 500000 + python build_thinking_rag_index.py --mode eval --eval-size 1000 + python build_thinking_rag_index.py --mode both --total 200000 --eval-size 500 +""" +import argparse +import json +import os +import re +import sys +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional, Tuple + +import numpy as np +import torch +import torch.nn.functional as F +from tqdm import tqdm + +# --------------------------------------------------------------------------- +# Reuse condenser prompts (single source of truth) +# --------------------------------------------------------------------------- +_HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(_HERE.parent / 'condenser')) +sys.path.insert(0, str(_HERE)) +from make_condenser_dataset import ( # noqa: E402 (after sys.path tweak) + COMPRESS_SYSTEM, + COMPRESS_USER, + FIXED_QUERY_NEED, + FIXED_QUERY_SKILL, +) +# Default dataset loader is the index-time corpus (broader retrieval profile); +# pass --dataset-module dataset_think to fall back to the training mix. +from dataset_index import get_dataset as _default_get_dataset # noqa: E402 + +_GET_DATASET = _default_get_dataset + +import twinkle # noqa: E402 +from twinkle import DeviceGroup, DeviceMesh, get_logger # noqa: E402 +from twinkle.data_format import SamplingParams as TwinkleSamplingParams # noqa: E402 +from twinkle.loss import InfonceLoss # noqa: E402 +from twinkle.model import TransformersModel # noqa: E402 +from twinkle.processor import InputProcessor # noqa: E402 +from twinkle.sampler import vLLMSampler # noqa: E402 +from twinkle.template import Qwen3_5Template # noqa: E402 +from twinkle.utils.parallel import PosixFileLock # noqa: E402 +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient # noqa: E402 + +logger = get_logger() + + +# =========================================================================== +# Config (most fields overridable via CLI / env) +# =========================================================================== + +EMBED_MODEL_ID = os.environ.get( + 'EMBED_MODEL_ID', + 'ms://twinkle-kit/Qwen3.5-4B-QA-emb', +) +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') + +# Twinkle device topology: TP=4 sampler on 0-3, DP=4 embedding on 4-7. +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 4)) +EMB_GPUS = int(os.environ.get('EMB_GPUS', 4)) +NUM_GPUS = SAMPLER_GPUS + EMB_GPUS + +# vLLM engine sizing. +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)) + +# Embedding sizing. +EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 8192)) + +SIM_THRESHOLD = float(os.environ.get('SIM_THRESHOLD', 0.65)) +MIN_TEXT_CHARS = int(os.environ.get('MIN_TEXT_CHARS', 256)) + +# OpenAI API fallback (used when vLLM truncates). +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_API_MODEL = os.environ.get('COMPRESS_API_MODEL', 'qwen3.7-max') + +# Source → coarse domain (for filtered eval). +DOMAIN_MAP = { + 'CodeX-2M-Thinking': 'code', + 'OpenThoughts3-1.2M': 'reasoning', + 'LIMO-v2': 'math', + 'Chinese-DeepSeek-R1-Distill-data-110k': 'reasoning_zh', + 'Opus-4.6-Reasoning-3000x-filtered': 'reasoning', + 'claude-opus-4.6-10000x': 'mixed', + 'angrygiraffe-claude-opus-4.6-4.7-reasoning-8.7k': 'mixed', +} + + +# =========================================================================== +# Small helpers +# =========================================================================== + +def _is_truncated_compression(text: str) -> bool: + """Detect structurally incomplete condenser output even when stop='stop'. + + The skeleton mandates a ``## Summary`` plus a ``## More`` bullet list whose + last line begins with ``-`` or ends with ``)`` (the ``(none)`` sentinel). + Anything short of that signals truncation; the API fallback is invoked. + """ + 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 + return False + + +def _strip_outer_codefence(text: str) -> str: + m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', text, re.DOTALL) + if m: + return m.group(1).strip() + return text.strip() + + +def _wrap_anchor(text: str) -> List[Dict[str, str]]: + """Anchor-side message wrapping (must match training).""" + return [ + {'role': 'user', 'content': text}, + {'role': 'assistant', 'content': 'Match the correct response here.'}, + ] + + +def _wrap_positive(text: str) -> List[Dict[str, str]]: + """Positive-side message wrapping (must match training).""" + return [ + {'role': 'user', 'content': 'Match the correct query here.'}, + {'role': 'assistant', 'content': text}, + ] + + +def _short(text: str, n: int = 96) -> str: + text = (text or '').replace('\n', ' ').strip() + return text[:n] + ('…' if len(text) > n else '') + + +def _detect_lang(text: str) -> str: + if not text: + return 'unknown' + cjk = sum(1 for ch in text[:512] if '\u4e00' <= ch <= '\u9fff') + return 'zh' if cjk >= 8 else 'en' + + +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)}, + ] + + +# =========================================================================== +# Twinkle component wrappers +# =========================================================================== + +def initialize_twinkle() -> Tuple[DeviceMesh, DeviceMesh]: + """Wire two device groups (sampler / emb_model) and return their meshes.""" + device_groups = [ + DeviceGroup( + name='sampler', + ranks=list(range(SAMPLER_GPUS)), + device_type='GPU', + gpus_per_worker=SAMPLER_GPUS, # TP=4 → one worker spans all 4 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, + ) + return sampler_mesh, emb_mesh + + +def build_sampler(sampler_mesh: DeviceMesh) -> vLLMSampler: + 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, + ) + return sampler + + +def build_emb_model(emb_mesh: DeviceMesh) -> Tuple[TransformersModel, Qwen3_5Template]: + model = TransformersModel( + model_id=EMBED_MODEL_ID, + device_mesh=emb_mesh, + remote_group='emb_model', + ) + model.set_processor(InputProcessor) + # InfonceLoss is required by the framework even though forward_only does + # not actually invoke it; matches the training-time configuration. + model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) + # Qwen3.5-specific subclass applies orphan- chat-template patches. + template = Qwen3_5Template( + model_id=EMBED_MODEL_ID, + max_length=EMBED_MAX_LENGTH, + truncation_strategy='delete', + enable_thinking=False, + ) + return model, template + + +# =========================================================================== +# Compression helpers (vLLMSampler) + API fallback +# =========================================================================== + +def _vllm_compress(sampler: vLLMSampler, texts: List[str], query_hint: str + ) -> List[Tuple[str, str]]: + """Compress ``texts`` via the sampler; return ``(decoded, stop_reason)``.""" + if not texts: + return [] + prompts = [{'messages': _build_compress_messages(t, query_hint)} for t in texts] + params = TwinkleSamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, + temperature=COMPRESS_TEMPERATURE, + top_p=COMPRESS_TOP_P, + num_samples=1, + ) + responses = sampler.sample(prompts, params) + results: List[Tuple[str, str]] = [] + for resp in responses: + seq = resp.sequences[0] if resp and resp.sequences else None + if seq is None: + results.append(('', 'error')) + continue + text = seq.decoded or '' + # Strip any leaked chat-template special tokens like ``<|im_end|>``. + text = re.sub(r'<\|[^|]+\|>', '', text).rstrip() + text = _strip_outer_codefence(text) + results.append((text, seq.stop_reason or 'stop')) + return results + + +def _api_compress(api: OpenAIClient, messages: List[Dict[str, str]]) -> Optional[str]: + sp = TwinkleSamplingParams(temperature=COMPRESS_TEMPERATURE, max_tokens=CONDENSE_MAX_TOKENS) + try: + reply = api({'messages': messages}, sp, extra_body={'enable_thinking': False}) + except Exception as exc: # noqa: BLE001 — broad catch is intentional + sys.stderr.write(f'[api_fallback] error: {exc}\n') + return None + content = (reply.get('content') or '').strip() + if not content: + return None + return _strip_outer_codefence(content) + + +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.""" + pairs = _vllm_compress(sampler, texts, query_hint) + results: List[Optional[str]] = [] + for (text, stop), src_text in zip(pairs, texts): + if stop != 'length' and not _is_truncated_compression(text): + results.append(text) + continue + if api is None: + results.append(None) + 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) + else: + results.append(api_text) + return results + + +# =========================================================================== +# Embedding helpers (TransformersModel.forward_only(task='embedding')) +# =========================================================================== + +def _build_features(template: Qwen3_5Template, texts: List[str], role: str + ) -> List[Dict[str, Any]]: + """Wrap each text into the role-specific anchor / positive feature dict.""" + features: List[Dict[str, Any]] = [] + for text in texts: + if not text or not text.strip(): + # Pad with a single space so positional alignment holds against + # the input list — the caller filters out empty-text rows upstream. + text = ' ' + if role == 'anchor': + feat = template.encode({'messages': _wrap_anchor(text)}) + feat['labels'] = [1] + else: + feat = template.encode({'messages': _wrap_positive(text)}) + feat['labels'] = [0] + features.append(feat) + return features + + +def get_embeddings(model: TransformersModel, template: Qwen3_5Template, + texts: List[str], role: str) -> np.ndarray: + """Return ``[N, H]`` float32 L2-normalised embeddings for ``texts``. + + Inputs are padded up to a multiple of ``EMB_GPUS`` and sliced back to the + original ``N``: the dispatch layer (``_dispatch_args``) starves any rank + whose chunk lands beyond ``len(texts)``, so a single forward of fewer than + ``EMB_GPUS`` items (e.g. the probe) would otherwise raise + ``Batch too small for {EMB_GPUS} workers``. + """ + 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 = _build_features(template, padded, role) + 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 _probe_hidden_size(model: TransformersModel, template: Qwen3_5Template) -> int: + """One-shot warmup forward to read out the embedding dimension.""" + emb = get_embeddings(model, template, ['probe'], role='anchor') + if emb.ndim != 2 or emb.shape[0] == 0: + raise RuntimeError(f'unexpected embedding shape from probe: {emb.shape}') + return int(emb.shape[1]) + + +# =========================================================================== +# LanceDB I/O +# =========================================================================== + +def _make_arrow_schema(hidden_size: int): + import pyarrow as pa + return pa.schema([ + pa.field('id', pa.string()), + pa.field('vector', pa.list_(pa.float32(), hidden_size)), + pa.field('thinking_raw', pa.string()), + pa.field('query_raw', pa.string()), + pa.field('cot_compressed', pa.string()), + pa.field('query_compressed', pa.string()), + pa.field('source', pa.string()), + pa.field('domain', pa.string()), + pa.field('language', pa.string()), + pa.field('sim', pa.float32()), + ]) + + +def _open_or_create_table(db_path: str, table_name: str, hidden_size: int, + mode: str): + """Open an existing table for append/eval, or create a fresh one.""" + import lancedb + db = lancedb.connect(db_path) + schema = _make_arrow_schema(hidden_size) + if table_name in db.table_names(): + if mode == 'overwrite': + db.drop_table(table_name) + tbl = db.create_table(table_name, schema=schema, mode='overwrite') + else: + tbl = db.open_table(table_name) + else: + tbl = db.create_table(table_name, schema=schema, mode='create') + return db, tbl + + +def _existing_ids(table) -> set: + try: + col = table.to_pandas(columns=['id']) + return set(col['id'].astype(str).tolist()) + except Exception: # noqa: BLE001 + return set() + + +# =========================================================================== +# Build pipeline +# =========================================================================== + +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) + 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' + + (f' → yielding first {cap}\n' if cap < n_full else '\n')) + for i, row in enumerate(ds): + if i >= cap: + break + yield row + + +def _extract_query_cot(row: Dict[str, Any]) -> Tuple[str, str]: + user_query, cot = '', '' + for m in row.get('messages') or []: + if not isinstance(m, dict): + continue + role = m.get('role') or '' + if role == 'user' and not user_query: + user_query = (m.get('content') or '').strip() + elif role == 'assistant': + cot = (m.get('reasoning_content') or '').strip() + break + return user_query, cot + + +def _log_miss(misses_path: str, lock: PosixFileLock, record: Dict[str, Any]) -> None: + line = json.dumps(record, ensure_ascii=False, default=str) + '\n' + with lock: + with open(misses_path, 'a', encoding='utf-8') as fh: + fh.write(line) + + +def build_index(args: argparse.Namespace, + sampler: vLLMSampler, + emb_model: TransformersModel, + emb_template: Qwen3_5Template, + api: Optional[OpenAIClient]) -> None: + # ---- Probe embedding dimension ----------------------------------------- + sys.stderr.write('[build] probing embedding hidden size...\n') + hidden_size = _probe_hidden_size(emb_model, emb_template) + sys.stderr.write(f'[build] hidden_size={hidden_size}\n') + + # ---- LanceDB ------------------------------------------------------------ + db, tbl = _open_or_create_table( + args.db_path, args.table, hidden_size, + mode='overwrite' if args.overwrite else 'append', + ) + indexed = _existing_ids(tbl) if not args.overwrite else set() + sys.stderr.write(f'[build] table "{args.table}" — {len(indexed)} existing rows.\n') + + misses_path = args.misses_log or (str(Path(args.db_path) / f'{args.table}.misses.jsonl')) + Path(misses_path).parent.mkdir(parents=True, exist_ok=True) + misses_lock = PosixFileLock(misses_path + '.lock') + + # ---- 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) + + batch: List[Dict[str, Any]] = [] + + def _flush(rows: List[Dict[str, Any]]) -> None: + nonlocal n_kept, n_dropped_compress, n_dropped_sim + if not rows: + return + # Phase 1 — compress query (FIXED_QUERY_NEED) and cot (FIXED_QUERY_SKILL). + q_compressed = _resolve_compressed( + sampler, api, [r['query_raw'] for r in rows], FIXED_QUERY_NEED) + c_compressed = _resolve_compressed( + sampler, api, [r['cot_raw'] for r in rows], FIXED_QUERY_SKILL) + kept_rows: List[Dict[str, Any]] = [] + for r, q_cmp, c_cmp in zip(rows, q_compressed, c_compressed): + if not q_cmp or not c_cmp: + 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), + 'cot_raw_head': _short(r['cot_raw'], 200), + }) + continue + r['query_compressed'] = q_cmp + r['cot_compressed'] = c_cmp + kept_rows.append(r) + 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' + print(f'[{tag} sim={sim_val:.4f}] {r["source"][:24]} ' + 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 + _log_miss(misses_path, misses_lock, { + 'id': r['id'], 'source': r['source'], 'reason': 'sim_low', + 'sim': float(sim_val), + 'query_raw': r['query_raw'], + 'cot_raw': r['cot_raw'], + 'query_compressed': r['query_compressed'], + 'cot_compressed': r['cot_compressed'], + }) + continue + to_insert.append({ + 'id': r['id'], + 'vector': positive_emb[idx].tolist(), + 'thinking_raw': r['cot_raw'], + 'query_raw': r['query_raw'], + 'cot_compressed': r['cot_compressed'], + 'query_compressed': r['query_compressed'], + 'source': r['source'], + 'domain': DOMAIN_MAP.get(r['source'], 'mixed'), + 'language': _detect_lang(r['cot_raw']), + 'sim': float(sim_val), + }) + if to_insert: + tbl.add(to_insert) + n_kept += len(to_insert) + indexed.update(r['id'] for r in to_insert) + + try: + 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: + break + rid = row.get('id') or '' + if not rid: + 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: + n_dropped_short += 1 + continue + batch.append({ + 'id': rid, + 'source': row.get('source') or 'unknown', + 'query_raw': user_query, + 'cot_raw': cot, + }) + if len(batch) >= args.batch_size: + _flush(batch) + batch.clear() + pbar.set_postfix(kept=n_kept, sim_drop=n_dropped_sim, + cmp_drop=n_dropped_compress, refresh=False) + pbar.update(1) + if batch: + _flush(batch) + batch.clear() + finally: + pbar.close() + + 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') + + # ---- Build vector index for fast retrieval ------------------------------ + if n_kept >= 64 and not args.skip_index: + sys.stderr.write('[build] creating IVF_PQ index (metric=dot)...\n') + n_partitions = max(8, min(256, n_kept // 1000 + 1)) + try: + tbl.create_index( + metric='dot', + vector_column_name='vector', + num_partitions=n_partitions, + num_sub_vectors=16, + index_type='IVF_PQ', + replace=True, + ) + except Exception as exc: # noqa: BLE001 + sys.stderr.write(f'[build] index build failed: {exc} ' + '(table is still queryable via brute-force scan)\n') + sys.stderr.write(f'[build] done. table rows={tbl.count_rows()}\n') + + +# =========================================================================== +# Eval pipeline (self-recall on indexed rows) +# =========================================================================== + +def eval_recall(args: argparse.Namespace, + sampler: vLLMSampler, + emb_model: TransformersModel, + emb_template: Qwen3_5Template, + api: Optional[OpenAIClient]) -> None: + """Probe each gold query against the index; report recall@k. + + Self-recall semantics: only rows whose ``id`` is already present in the + index are probed. The corresponding ``cot``-keyed vector must be retrieved + by encoding the **raw user query** through the condenser → embedder + pipeline (anchor side). The match is correct iff the retrieved row's + ``id`` equals the probe row's ``id``. + """ + import lancedb + db = lancedb.connect(args.db_path) + if args.table not in db.table_names(): + raise SystemExit(f'[eval] table "{args.table}" does not exist in {args.db_path}') + tbl = db.open_table(args.table) + indexed_ids = _existing_ids(tbl) + sys.stderr.write(f'[eval] table rows={tbl.count_rows()} indexed_ids={len(indexed_ids)}\n') + if not indexed_ids: + sys.stderr.write('[eval] empty index — nothing to evaluate.\n') + return + + 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] = {} + probed = 0 + + pbar = tqdm(desc='eval', unit='probe', dynamic_ncols=True) + batch_rows: List[Dict[str, Any]] = [] + + def _flush(rows: List[Dict[str, Any]]) -> None: + nonlocal probed + if not rows: + return + compressed = _resolve_compressed( + sampler, api, [r['query_raw'] for r in rows], FIXED_QUERY_NEED) + useful = [(r, c) for r, c in zip(rows, compressed) if c] + if not useful: + return + anchor_emb = get_embeddings( + emb_model, emb_template, [c for _, c in useful], role='anchor') + for (r, _), vec in zip(useful, anchor_emb): + res = ( + tbl.search(vec.astype(np.float32).tolist()) + .metric('dot') + .limit(max(ks)) + .select(['id', 'source']) + .to_list() + ) + hit_ids = [item['id'] for item in res] + try: + rank = hit_ids.index(r['id']) + except ValueError: + rank = -1 + for k in ks: + if 0 <= rank < k: + hits[k] += 1 + per_source_hits.setdefault(r['source'], {kk: 0 for kk in ks})[k] += 1 + per_source_total[r['source']] = per_source_total.get(r['source'], 0) + 1 + per_source_hits.setdefault(r['source'], {kk: 0 for kk in ks}) + probed += 1 + pbar.update(len(useful)) + + try: + for row in _stream_corpus(total=args.total, load_from_cache_file=not args.no_cache, + max_rows=args.max_rows): + if probed + len(batch_rows) >= args.eval_size: + break + rid = row.get('id') or '' + if not rid or rid not in indexed_ids: + continue + user_query, _ = _extract_query_cot(row) + if not user_query or len(user_query) < MIN_TEXT_CHARS: + continue + batch_rows.append({ + 'id': rid, + 'source': row.get('source') or 'unknown', + 'query_raw': user_query, + }) + if len(batch_rows) >= args.batch_size: + _flush(batch_rows) + batch_rows.clear() + if batch_rows: + _flush(batch_rows) + finally: + pbar.close() + + if probed == 0: + sys.stderr.write( + '[eval] no probed rows — index empty, queries too short, or ' + 'corpus exhausted before eval-size?\n') + return + + print('\n=== Recall @ k (self-recall, gold present in index) ===') + print(f'probed = {probed}') + for k in ks: + print(f' recall@{k:<3} = {hits[k]/probed:.4f} ({hits[k]}/{probed})') + + print('\n=== Per-source recall@10 ===') + for src in sorted(per_source_total): + tot = per_source_total[src] + h10 = per_source_hits.get(src, {}).get(10, 0) + print(f' {src:<48s} {h10/tot:.4f} ({h10}/{tot})') + + +# =========================================================================== +# CLI +# =========================================================================== + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--mode', choices=['build', 'eval', 'both'], default='build') + p.add_argument('--db-path', default='./output/thinking_rag/lance.db', + help='LanceDB on-disk directory (persisted across runs).') + p.add_argument('--table', default='thinking_traces', + 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('--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('--no-cache', action='store_true', + help='Disable load_from_cache_file in dataset_think.get_dataset.') + p.add_argument('--overwrite', action='store_true', + help='Drop the table before build and start fresh.') + p.add_argument('--skip-index', action='store_true', + help='Skip IVF_PQ index build at the end (debug).') + p.add_argument('--misses-log', default='', + help='Path for filtered-row JSONL log (defaults to /.misses.jsonl).') + + # eval-only + p.add_argument('--eval-size', type=int, default=500, + help='Number of probes for self-recall evaluation.') + p.add_argument('--top-k', type=int, default=10, + help='Largest k to report. Smaller ks (1, 5) are always reported.') + + return p.parse_args() + + +def main() -> None: + args = parse_args() + Path(args.db_path).mkdir(parents=True, exist_ok=True) + + global _GET_DATASET + if args.dataset_module == 'dataset_think': + from dataset_think import get_dataset as _swap + _GET_DATASET = _swap + sys.stderr.write(f'[main] dataset loader: {args.dataset_module}\n') + + # Build/eval both depend on the same Twinkle stack — initialize once. + sampler_mesh, emb_mesh = initialize_twinkle() + sys.stderr.write(f'[main] twinkle initialized: ' + f'sampler ranks 0-{SAMPLER_GPUS - 1} (TP={SAMPLER_GPUS}), ' + f'emb_model ranks {SAMPLER_GPUS}-{NUM_GPUS - 1} (DP={EMB_GPUS}).\n') + + sys.stderr.write('[main] starting vLLM condenser sampler...\n') + sampler = build_sampler(sampler_mesh) + sys.stderr.write('[main] starting embedding TransformersModel...\n') + emb_model, emb_template = build_emb_model(emb_mesh) + + api: Optional[OpenAIClient] = None + if COMPRESS_API_KEY: + api = OpenAIClient( + model=COMPRESS_API_MODEL, + api_key=COMPRESS_API_KEY, + base_url=COMPRESS_BASE_URL, + ) + else: + sys.stderr.write( + '[main] WARNING: COMPRESS_API_KEY unset — truncated rows will be dropped.\n') + + if args.mode in ('build', 'both'): + build_index(args, sampler, emb_model, emb_template, api) + if args.mode in ('eval', 'both'): + eval_recall(args, sampler, emb_model, emb_template, api) + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/dataset_index.py b/cookbook/exp/embedding/dataset_index.py new file mode 100644 index 00000000..7d2905a5 --- /dev/null +++ b/cookbook/exp/embedding/dataset_index.py @@ -0,0 +1,718 @@ +"""RAG-index corpus loader — abstract reasoning skills + textbook-style methods. + +Distinct from training-time ``dataset_think.py``. Optimizes for **abstraction +density**, not raw coverage: every row should encode a transferable method, +theorem, or solution pattern that downstream queries can retrieve as a +"use-when-X-do-Y" recipe. + +Single-table design (``thinking_traces``); EMBED_QUERY_COT condense step in +``build_thinking_rag_index`` homogenizes thinking-style and textbook-style +content into the same retrieval form, so dual-table is unnecessary. The +``source`` field carries the original dataset name for eval-time +domain-bucket diagnostics. + +Output schema matches ``dataset_think.get_dataset()``: ``{id, source, messages}`` +with ``messages[1].reasoning_content`` carrying the CoT. + +Mix (≈3.6M rows base, 10 datasets): + Math thinking 23% — OpenMathReasoning + OpenR1-Math-220k + s1K-1.1 + Code thinking 19% — OpenCodeReasoning-2 + codeforces-cots + Cross-domain R1 39% — Bespoke-Stratos + dolphin-r1 + reasoning-v1-20m + + natural_reasoning + Textbook synth 17% — cosmopedia v1 (auto_math_text, chunked by H2) + Olympiad solutions <1% — Omni-MATH + +Dropped: camel-ai/{physics,chemistry,biology} (zip-only, no parquet/jsonl) and +swift/stack-exchange-paired (dataset_infos.json/data layout mismatch); the +textbook-density gap is covered by a larger cosmopedia slice. + +Textbook processors synthesize a question from the chapter heading and place +the explanatory body into the ``cot`` field — embedding+condense reads +``query | cot`` so the textbook prose becomes a retrievable method. + +Field extraction is defensive: each processor tries multiple plausible column +names and silently drops rows that miss a usable signal. Inspect +``dropped_index.jsonl`` after the first run to verify field-name guesses. +""" +import re +from typing import Any, Dict, List, Optional + +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.preprocessor import Preprocessor + +from dataset_think import _THINK_RE, _hash_id, _register, ToMessagesProcessor + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +# Sky-T1 / Bespoke-Stratos custom markers (used in place of ). +_BOT_RE = re.compile( + r'<\|begin_of_thought\|>(.*?)<\|end_of_thought\|>', re.DOTALL) +_BOS_RE = re.compile( + r'<\|begin_of_solution\|>(.*?)<\|end_of_solution\|>', re.DOTALL) + +# H2 heading split for cosmopedia-style markdown chunks. +_H2_RE = re.compile(r'^##\s+(.+?)\s*$', re.MULTILINE) + + +def _split_think(text: str) -> tuple: + """Return ``(cot, response)``; cot empty if no ```` block found.""" + if not text: + return '', '' + m = _THINK_RE.search(text) + if not m: + return '', text.strip() + return m.group(1).strip(), text[m.end():].strip() + + +def _split_sky_t1(text: str) -> tuple: + """Return ``(cot, response)`` for Sky-T1 / Bespoke-Stratos marker format.""" + if not text: + return '', '' + bot = _BOT_RE.search(text) + bos = _BOS_RE.search(text) + cot = bot.group(1).strip() if bot else '' + sol = bos.group(1).strip() if bos else '' + return cot, sol + + +def _from_messages(messages: Any) -> tuple: + """Pull (first_user, first_assistant) from OpenAI/ShareGPT-style list.""" + if not isinstance(messages, list): + return '', '' + query, assistant = '', '' + for msg in messages: + if not isinstance(msg, dict): + continue + role = msg.get('role') or msg.get('from') or '' + content = msg.get('content') or msg.get('value') or '' + if not isinstance(content, str): + continue + if role in ('user', 'human') and not query: + query = content.strip() + elif role in ('assistant', 'gpt') and not assistant: + assistant = content.strip() + break + return query, assistant + + +def _chunk_by_h2(text: str, min_chars: int = 200, max_chars: int = 6000): + """Split markdown text on ``## `` headings; yield ``(title, body)`` pairs.""" + if not text: + return + matches = list(_H2_RE.finditer(text)) + if not matches: + head = text.strip()[:80].splitlines()[0] if text.strip() else '' + body = text.strip() + if head and min_chars <= len(body) <= max_chars: + yield head, body + return + for i, m in enumerate(matches): + title = m.group(1).strip() + start = m.end() + end = matches[i + 1].start() if i + 1 < len(matches) else len(text) + body = text[start:end].strip() + if min_chars <= len(body) <= max_chars and title: + yield title, body + + +# =========================================================================== +# Math thinking +# =========================================================================== + +OPEN_MATH_REASONING_REPO = 'ms://AI-ModelScope/OpenMathReasoning' + + +class OpenMathReasoningProcessor(Preprocessor): + """OpenMathReasoning → ``{id, source, query, cot, response}``. + + Schema: ``problem``, ``generated_solution`` (R1 trace with ````), + ``expected_answer``. The ``cot`` *split* (not column) is the long-CoT + portion — TIR/genselect/additional_problems sit in sibling splits and + are filtered at load time, not row-level. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + query = (row.get('problem') or row.get('question') or '').strip() + assistant = (row.get('generated_solution') or row.get('solution') + or row.get('output') or '').strip() + if not query or not assistant: + continue + cot, response = _split_think(assistant) + if not cot: + continue + if not response: + response = (row.get('expected_answer') or row.get('answer') or '').strip() + if not response: + continue + out.append({ + 'id': _hash_id('open_math_reasoning', f'{query}\n{response}'), + 'source': 'OpenMathReasoning', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +OPEN_R1_MATH_REPO = 'ms://open-r1/OpenR1-Math-220k' + + +class OpenR1MathProcessor(Preprocessor): + """OpenR1-Math-220k → ``{id, source, query, cot, response}``. + + Schema: ``problem``, ``solution``, ``answer``, ``generations`` (list of + R1 traces), ``correctness_math_verify`` (parallel bool list). Pick the + first generation whose math-verify passed; fall back to ``solution``. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + query = (row.get('problem') or row.get('question') or '').strip() + if not query: + continue + assistant = '' + gens = row.get('generations') + verifies = row.get('correctness_math_verify') + if isinstance(gens, list): + if isinstance(verifies, list) and len(verifies) == len(gens): + for g, v in zip(gens, verifies): + if v and isinstance(g, str) and g.strip(): + assistant = g.strip() + break + if not assistant: + for g in gens: + if isinstance(g, str) and g.strip(): + assistant = g.strip() + break + if not assistant: + assistant = (row.get('solution') or '').strip() + if not assistant: + continue + cot, response = _split_think(assistant) + if not cot: + continue + if not response: + response = (row.get('answer') or '').strip() + if not response: + continue + out.append({ + 'id': _hash_id('open_r1_math', f'{query}\n{response}'), + 'source': 'OpenR1-Math-220k', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +S1K_REPO = 'ms://simplescaling/s1K-1.1' + + +class S1KProcessor(Preprocessor): + """s1K-1.1 → ``{id, source, query, cot, response}``. + + Schema: ``question`` + ``deepseek_thinking_trajectory`` (or + ``thinking_trajectories`` legacy) + ``deepseek_attempt`` (final answer). + Hand-curated peak-abstraction set, kept whole. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + query = (row.get('question') or row.get('problem') or '').strip() + thinking = (row.get('deepseek_thinking_trajectory') + or row.get('thinking_trajectories') + or row.get('thinking') or '') + if isinstance(thinking, list): + thinking = '\n\n'.join(t for t in thinking if isinstance(t, str)) + cot = (thinking or '').strip() + response = (row.get('deepseek_attempt') or row.get('attempt') + or row.get('answer') or row.get('solution') or '').strip() + if not query or not cot or not response: + continue + out.append({ + 'id': _hash_id('s1k', f'{query}\n{response}'), + 'source': 's1K-1.1', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +# =========================================================================== +# Code thinking +# =========================================================================== + +OPEN_CODE_REASONING_REPO = 'ms://nv-community/OpenCodeReasoning-2' + + +class OpenCodeReasoning2Processor(Preprocessor): + """OpenCodeReasoning-2 → ``{id, source, query, cot, response}``. + + Schema: ``input``/``problem``, plus per-model R1-style trace columns + (``r1_generation``, ``qwq_generation``, etc.). Prefer the ``r1`` trace; + fall back to ``solution``. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + query = (row.get('input') or row.get('problem') + or row.get('question') or '').strip() + # OCR-2 'python' split ships dirty rows where question is literally '-'; + # the real prompt is buried in r1_generation and not recoverable here. + if not query or query == '-': + continue + assistant = (row.get('r1_generation') or row.get('reasoning_content') + or row.get('solution') or row.get('output') or '').strip() + if not assistant: + continue + cot, response = _split_think(assistant) + if not cot: + continue + if not response: + response = (row.get('expected_solution') or row.get('answer') or '').strip() + if not response: + continue + out.append({ + 'id': _hash_id('opencode_reasoning2', f'{query}\n{response}'), + 'source': 'OpenCodeReasoning-2', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +CODEFORCES_COTS_REPO = 'ms://open-r1/codeforces-cots' + + +class CodeforcesCotsProcessor(Preprocessor): + """codeforces-cots → ``{id, source, query, cot, response}``. + + Schema: ``description``/``problem``, ``generation``/``solution`` (R1 + trace with ```` + final code). Algorithmic patterns at high + abstraction density. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + query = (row.get('description') or row.get('problem') + or row.get('input') or row.get('question') or '').strip() + assistant = (row.get('generation') or row.get('solution') + or row.get('output') or '').strip() + if not query or not assistant: + continue + cot, response = _split_think(assistant) + if not cot or not response: + continue + out.append({ + 'id': _hash_id('codeforces_cots', f'{query}\n{response}'), + 'source': 'codeforces-cots', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +# =========================================================================== +# Cross-domain R1 +# =========================================================================== + +BESPOKE_STRATOS_REPO = 'ms://bespokelabs/Bespoke-Stratos-17k' + + +class BespokeStratosProcessor(Preprocessor): + """Bespoke-Stratos-17k → ``{id, source, query, cot, response}``. + + Schema: ``conversations`` (ShareGPT). Assistant content uses Sky-T1 + markers ``<|begin_of_thought|>...<|end_of_thought|>`` then + ``<|begin_of_solution|>...<|end_of_solution|>``. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + query, assistant = _from_messages( + row.get('conversations') or row.get('messages')) + if not query or not assistant: + continue + cot, response = _split_sky_t1(assistant) + if not cot: + cot, response = _split_think(assistant) + if not cot or not response: + continue + out.append({ + 'id': _hash_id('bespoke_stratos', f'{query}\n{response}'), + 'source': 'Bespoke-Stratos-17k', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +DOLPHIN_R1_REPO = 'ms://AI-ModelScope/dolphin-r1' + + +class DolphinR1Processor(Preprocessor): + """dolphin-r1 → ``{id, source, query, cot, response}``. + + Schema (reasoning-deepseek subset): ``messages=[system, user]`` (no + assistant turn) + flat ``reasoning`` (CoT) + ``answer`` (final response) + + ``model``. Pull the user turn as query, ``reasoning``/``answer`` as + cot/response. Fallback to embedded ```` for legacy rows. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + msgs = row.get('messages') or row.get('conversations') + query = '' + if isinstance(msgs, list): + for msg in msgs: + if not isinstance(msg, dict): + continue + role = msg.get('role') or msg.get('from') or '' + content = msg.get('content') or msg.get('value') or '' + if role in ('user', 'human') and isinstance(content, str): + query = content.strip() + cot = (row.get('reasoning') or row.get('reasoning_content') or '').strip() + response = (row.get('answer') or '').strip() + if (not cot or not response) and isinstance(msgs, list): + _, assistant = _from_messages(msgs) + if assistant: + c2, r2 = _split_think(assistant) + if c2: + cot = cot or c2 + response = response or r2 or assistant + if not query or not cot or not response: + continue + out.append({ + 'id': _hash_id('dolphin_r1', f'{query}\n{response}'), + 'source': 'dolphin-r1', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +GLAIVE_REASONING_REPO = 'ms://glaiveai/reasoning-v1-20m' + + +class GlaiveReasoningProcessor(Preprocessor): + """reasoning-v1-20m → ``{id, source, query, cot, response}``. + + Schema: ``prompt``, ``response`` (R1 trace with ```` + answer). + Largest cross-domain corpus in the mix; downsample aggressively. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + query = (row.get('prompt') or row.get('question') + or row.get('input') or '').strip() + assistant = (row.get('response') or row.get('output') + or row.get('answer') or '').strip() + if not query or not assistant: + continue + cot, response = _split_think(assistant) + if not cot or not response: + continue + out.append({ + 'id': _hash_id('glaive_reasoning', f'{query}\n{response}'), + 'source': 'reasoning-v1-20m', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +NATURAL_REASONING_REPO = 'ms://facebook/natural_reasoning' + + +class NaturalReasoningProcessor(Preprocessor): + """natural_reasoning → ``{id, source, query, cot, response}``. + + Schema: ``question`` + ``reference_answer`` + ``responses=[{response_model, + response}]``. The ``response`` field itself is the step-by-step CoT + (``## Step 1...## Step 2...``); there is no separate ``reasoning`` key. + Map ``responses[i].response`` → cot, ``reference_answer`` → response. + Rows with empty ``reference_answer`` (~18% per README) are dropped. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + query = (row.get('question') or '').strip() + if not query: + continue + cot = '' + responses = row.get('responses') + if isinstance(responses, list): + for r in responses: + if not isinstance(r, dict): + continue + txt = (r.get('response') or r.get('reasoning') + or r.get('thinking') or r.get('answer') or '').strip() + if txt: + cot = txt + break + if not cot: + cot = (row.get('reasoning') or row.get('thinking') + or row.get('response') or '').strip() + response = (row.get('reference_answer') or row.get('answer') or '').strip() + if not cot or not response: + continue + out.append({ + 'id': _hash_id('natural_reasoning', f'{query}\n{response}'), + 'source': 'natural_reasoning', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +# =========================================================================== +# Textbook-style — synthesize query from chapter heading; body → cot +# =========================================================================== + +COSMOPEDIA_REPO = 'ms://HuggingFaceTB/cosmopedia' + +class CosmopediaProcessor(Preprocessor): + """cosmopedia v1 → ``{id, source, query, cot, response}``. + + Schema: ``prompt`` (writing instruction), ``text`` (full chapter body), + ``format``/``audience``/``seed_data``. The subset is selected at load + time (``subset_name='auto_math_text'`` — densest math-textbook slice); + H2 chunking inside each row yields synthetic queries + (``Explain {heading}``) with the body placed into ``cot``. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + text = (row.get('text') or row.get('content') or '').strip() + if not text: + continue + for title, body in _chunk_by_h2(text): + # Heading-only "Explain: X" was 1-2 tokens and impossible to align + # with full-section cot. Promote the section's lead paragraph into + # the query so anchor carries real semantic content. + parts = body.split('\n\n', 1) + first_para = parts[0].strip() + rest = parts[1].strip() if len(parts) > 1 else '' + if len(first_para) < 256 or len(rest) < 256: + continue + query = f'{title}\n\n{first_para}' if title else first_para + out.append({ + 'id': _hash_id('cosmopedia', f'{title}\n{first_para[:200]}'), + 'source': 'cosmopedia-v1', + 'query': query, + 'cot': rest, + 'response': '', + }) + return self.map_row_to_col(out) + + +OMNI_MATH_REPO = 'ms://AI-ModelScope/Omni-MATH' + + +class OmniMathProcessor(Preprocessor): + """Omni-MATH → ``{id, source, query, cot, response}``. + + Schema: ``problem``, ``solution`` (full proof), ``answer``, ``domain``, + ``difficulty``. Olympiad-grade derivations — solution body → cot, + answer → response. + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + query = (row.get('problem') or row.get('question') or '').strip() + solution = (row.get('solution') or '').strip() + answer = (row.get('answer') or row.get('expected_answer') or '').strip() + if not query or not solution: + continue + out.append({ + 'id': _hash_id('omni_math', f'{query}\n{solution[:200]}'), + 'source': 'Omni-MATH', + 'query': query, + 'cot': solution, + 'response': answer, + }) + return self.map_row_to_col(out) + + +# =========================================================================== +# Mix configuration — base sizes target ≈3.6M total rows +# =========================================================================== + +_BASE_SIZES = { + 'open_math_reasoning': 600_000, + 'open_r1_math': 220_000, + 's1k': 1_000, + 'opencode_reasoning2': 500_000, + 'codeforces_cots': 200_000, + 'bespoke_stratos': 17_000, + 'dolphin_r1': 400_000, + 'glaive_reasoning': 800_000, + 'natural_reasoning': 200_000, + 'cosmopedia': 700_000, + 'omni_math': 4_000, +} + + +def _scaled_sizes(total: Optional[int]) -> Dict[str, int]: + if total is None or total <= 0: + return dict(_BASE_SIZES) + scale = total / sum(_BASE_SIZES.values()) + return {k: max(1, int(round(v * scale))) for k, v in _BASE_SIZES.items()} + + +def _build_dataset(total: Optional[int] = None, + load_from_cache_file: bool = True) -> Dataset: + sizes = _scaled_sizes(total) + dataset = Dataset() + + _register(dataset, OpenMathReasoningProcessor, + DatasetMeta(dataset_id=OPEN_MATH_REASONING_REPO, split='cot', + data_slice=range(sizes['open_math_reasoning'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, OpenR1MathProcessor, + DatasetMeta(dataset_id=OPEN_R1_MATH_REPO, split='train', + data_slice=range(sizes['open_r1_math'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, S1KProcessor, + DatasetMeta(dataset_id=S1K_REPO, split='train'), + load_from_cache_file=load_from_cache_file) + + _register(dataset, OpenCodeReasoning2Processor, + DatasetMeta(dataset_id=OPEN_CODE_REASONING_REPO, + subset_name='train', split='python', + data_slice=range(sizes['opencode_reasoning2'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, CodeforcesCotsProcessor, + DatasetMeta(dataset_id=CODEFORCES_COTS_REPO, + subset_name='solutions_w_editorials_decontaminated', + split='train', + data_slice=range(sizes['codeforces_cots'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, BespokeStratosProcessor, + DatasetMeta(dataset_id=BESPOKE_STRATOS_REPO, split='train'), + load_from_cache_file=load_from_cache_file) + + _register(dataset, DolphinR1Processor, + DatasetMeta(dataset_id=DOLPHIN_R1_REPO, + subset_name='reasoning-deepseek', split='train', + data_slice=range(sizes['dolphin_r1'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, GlaiveReasoningProcessor, + DatasetMeta(dataset_id=GLAIVE_REASONING_REPO, split='train', + data_slice=range(sizes['glaive_reasoning'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, NaturalReasoningProcessor, + DatasetMeta(dataset_id=NATURAL_REASONING_REPO, split='train', + data_slice=range(sizes['natural_reasoning'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, CosmopediaProcessor, + DatasetMeta(dataset_id=COSMOPEDIA_REPO, + subset_name='auto_math_text', split='train', + data_slice=range(sizes['cosmopedia'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, OmniMathProcessor, + DatasetMeta(dataset_id=OMNI_MATH_REPO, split='test'), + load_from_cache_file=load_from_cache_file) + + dataset.mix_dataset(False) + # Mix is concatenated in registration order; shuffle so the streaming + # consumer sees all sources interleaved instead of 600k OpenMathReasoning + # rows before it ever reaches code/textbook splits. + dataset.dataset = dataset.dataset.shuffle(seed=42) + return dataset + + +def get_dataset(total: Optional[int] = None, + dropped_log: Optional[str] = None, + load_from_cache_file: bool = True) -> Dataset: + """Build, convert to messages, and quality-filter the RAG-index corpus. + + Mirrors ``dataset_think.get_dataset``: identical signature + output + schema so ``build_thinking_rag_index`` consumes both modules unchanged. + """ + from twinkle_agentic.preprocessor import ( + DeadLoopFilter, + FixUnicodeFilter, + HardFilter, + MessageSanityFilter, + QualityPreprocessor, + RefuseFilter, + RemoveRepeatSentencesFilter, + TokenNumFilter, + TokenSoupFilter, + ) + + dataset = _build_dataset(total=total, load_from_cache_file=load_from_cache_file) + # Drop trivially-short queries (e.g. one-line math problems, OmniMath stubs) + # before message conversion — anchor side needs enough tokens to embed meaningfully. + dataset.dataset = dataset.dataset.filter( + lambda x: len((x.get('query') or '').strip()) >= 100, + num_proc=32, load_from_cache_file=load_from_cache_file) + dataset.map(ToMessagesProcessor(), remove_columns=['query', 'cot', 'response'], + load_from_cache_file=load_from_cache_file) + qp = QualityPreprocessor( + pipeline=[ + HardFilter(), + RefuseFilter(), + DeadLoopFilter(), + TokenSoupFilter(), + MessageSanityFilter(min_turns=1, max_msg_chars=200000), + FixUnicodeFilter(), + RemoveRepeatSentencesFilter(), + TokenNumFilter(max_num=32768), + ], + dropped_log_path=dropped_log or '', + ) + dataset.map(qp, num_proc=32, load_from_cache_file=load_from_cache_file) + return dataset + + +if __name__ == '__main__': + import os + dropped_log = os.path.join(os.path.dirname(os.path.abspath(__file__)), + 'dropped_index.jsonl') + if os.path.exists(dropped_log): + os.remove(dropped_log) + dataset = get_dataset(load_from_cache_file=False) + print(len(dataset)) diff --git a/cookbook/exp/embedding/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py index 492e29aa..bed1dbe9 100644 --- a/cookbook/exp/embedding/train_embedding_full_ddp.py +++ b/cookbook/exp/embedding/train_embedding_full_ddp.py @@ -40,7 +40,8 @@ from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient sys.path.insert(0, str(Path(__file__).resolve().parent)) -from dataset_think import get_dataset # noqa: E402 +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() @@ -49,7 +50,7 @@ # 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') +MODEL_ID = os.environ.get('MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-QA-emb') TEMPLATE_NAME = 'Qwen3_5Template' # -- GPU placement (8 total) -------------------------------------------------- @@ -68,15 +69,19 @@ GRADIENT_ACCUMULATION_STEPS = 1 LOG_INTERVAL = 2 SAVE_INTERVAL = 4000 -NUM_EPOCHS = 2 +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] = 200_000 +INDEX_CAP: Optional[int] = 400_000 +MIX_SHUFFLE_SEED = 42 # -- Resume from checkpoint --------------------------------------------------- -RESUME_CHECKPOINT = os.environ.get( - 'RESUME_CHECKPOINT', - './output/embedding_lora_transformers/step_16000') -RESUME_STEP = int(os.environ.get('RESUME_STEP', 16000)) +# 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; @@ -355,10 +360,15 @@ def _build_compress_prompts(rows: List[Dict[str, Any]]) -> tuple: raw_pairs: List[tuple] = [] prompt_queries: List[str] = [] passthrough: List[Optional[str]] = [] + # Conservative char budget: 32768 max_length - 8192 gen - ~2k prompt overhead = ~22k tokens. + # At worst-case 1.5 chars/token (CJK), that's ~33k chars. Use 60k as safe English-mix cap. + _MAX_COT_CHARS = 60_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)) # Short query bypasses condenser to avoid skeleton-induced hallucination. @@ -521,7 +531,20 @@ def train(): twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups) # -------- Data ----------------------------------------------------------- - dataset = get_dataset(total=TOTAL_SAMPLES, load_from_cache_file=True) + 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)}') dataloader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, shuffle=True) total_forward_steps = len(dataloader) * NUM_EPOCHS optimizer_steps = total_forward_steps // GRADIENT_ACCUMULATION_STEPS @@ -608,15 +631,19 @@ def _sample_batch(raw_batch): """Compress via vLLM sampler; fall back to API on truncation.""" compress_prompts, valid_indices, raw_pairs, prompt_queries, passthrough = \ _build_compress_prompts(raw_batch) - if not compress_prompts: + if len(compress_prompts) < 4: return None # 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: - with retrainer.sampler_lock: - sampler_responses = condenser_sampler.sample(sampler_input, compress_params) + try: + with retrainer.sampler_lock: + sampler_responses = condenser_sampler.sample(sampler_input, compress_params) + except Exception as exc: + logger.warning(f'[sampler] encode overflow in batch, falling back to API: {exc}') + sampler_responses = [None] * len(sampler_input) else: sampler_responses = [] responses = [None] * len(compress_prompts) From a79fb3b4b62f368c3c2a5f2c4e99fd583465f02e Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Mon, 15 Jun 2026 22:00:24 +0800 Subject: [PATCH 02/26] fix --- .../exp/embedding/build_thinking_rag_index.py | 66 +++-- .../exp/embedding/train_embedding_full_ddp.py | 242 +++++++++--------- 2 files changed, 170 insertions(+), 138 deletions(-) diff --git a/cookbook/exp/embedding/build_thinking_rag_index.py b/cookbook/exp/embedding/build_thinking_rag_index.py index cf4d1dd7..0299aa48 100644 --- a/cookbook/exp/embedding/build_thinking_rag_index.py +++ b/cookbook/exp/embedding/build_thinking_rag_index.py @@ -2,9 +2,10 @@ Pipeline (per row, batched): 1. Load (user_query, reasoning_content) pairs from ``dataset_think.get_dataset``. - 2. Compress query with ``FIXED_QUERY_NEED`` and cot with ``FIXED_QUERY_SKILL`` - using a Twinkle ``vLLMSampler`` (TP=4 across GPUs 0-3). Reuses the prompt - suite from ``cookbook/exp/condenser/make_condenser_dataset.py``. + 2. Compress query with ``RAG_QUERY_HINT`` and cot with ``RAG_THINKING_HINT`` + (a symmetric Problem/Skill/Knowledge schema defined in this file) using a + Twinkle ``vLLMSampler`` (TP=4 across GPUs 0-3). Reuses the system/user + wrappers from ``cookbook/exp/condenser/make_condenser_dataset.py``. 3. On condenser truncation (``stop_reason='length'`` or skeleton-incomplete output), fall back to an external OpenAI-compatible API. 4. Encode the condensed pair via the trained embedding model — Twinkle @@ -55,8 +56,6 @@ from make_condenser_dataset import ( # noqa: E402 (after sys.path tweak) COMPRESS_SYSTEM, COMPRESS_USER, - FIXED_QUERY_NEED, - FIXED_QUERY_SKILL, ) # Default dataset loader is the index-time corpus (broader retrieval profile); # pass --dataset-module dataset_think to fall back to the training mix. @@ -84,7 +83,7 @@ EMBED_MODEL_ID = os.environ.get( 'EMBED_MODEL_ID', - 'ms://twinkle-kit/Qwen3.5-4B-QA-emb', + '/mnt/workspace/yzhao/tastelikefeet/Qwen3.5-4B-QA-emb-v2', ) CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') @@ -106,6 +105,31 @@ SIM_THRESHOLD = float(os.environ.get('SIM_THRESHOLD', 0.65)) MIN_TEXT_CHARS = int(os.environ.get('MIN_TEXT_CHARS', 256)) +# 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' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the specific pattern, never generic labels.') +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' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the specific pattern, never generic labels.') + # OpenAI API fallback (used when vLLM truncates). COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '') COMPRESS_BASE_URL = os.environ.get( @@ -128,12 +152,19 @@ # Small helpers # =========================================================================== -def _is_truncated_compression(text: str) -> bool: - """Detect structurally incomplete condenser output even when stop='stop'. +_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') +_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') + - The skeleton mandates a ``## Summary`` plus a ``## More`` bullet list whose - last line begins with ``-`` or ends with ``)`` (the ``(none)`` sentinel). - Anything short of that signals truncation; the API fallback is invoked. +def _is_truncated_compression(text: str) -> 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 + * 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. """ if not text or not text.strip(): return True @@ -145,6 +176,11 @@ def _is_truncated_compression(text: str) -> bool: last_line = after_more.splitlines()[-1].strip() if not (last_line.startswith('-') or last_line.endswith(')')): return True + 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 @@ -494,11 +530,11 @@ def _flush(rows: List[Dict[str, Any]]) -> None: nonlocal n_kept, n_dropped_compress, n_dropped_sim if not rows: return - # Phase 1 — compress query (FIXED_QUERY_NEED) and cot (FIXED_QUERY_SKILL). + # Phase 1 — compress query (RAG_QUERY_HINT) and cot (RAG_THINKING_HINT). q_compressed = _resolve_compressed( - sampler, api, [r['query_raw'] for r in rows], FIXED_QUERY_NEED) + sampler, api, [r['query_raw'] for r in rows], RAG_QUERY_HINT) c_compressed = _resolve_compressed( - sampler, api, [r['cot_raw'] for r in rows], FIXED_QUERY_SKILL) + sampler, api, [r['cot_raw'] for r in rows], RAG_THINKING_HINT) kept_rows: List[Dict[str, Any]] = [] for r, q_cmp, c_cmp in zip(rows, q_compressed, c_compressed): if not q_cmp or not c_cmp: @@ -655,7 +691,7 @@ def _flush(rows: List[Dict[str, Any]]) -> None: if not rows: return compressed = _resolve_compressed( - sampler, api, [r['query_raw'] for r in rows], FIXED_QUERY_NEED) + sampler, api, [r['query_raw'] for r in rows], RAG_QUERY_HINT) useful = [(r, c) for r, c in zip(rows, compressed) if c] if not useful: return diff --git a/cookbook/exp/embedding/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py index bed1dbe9..9ac925b6 100644 --- a/cookbook/exp/embedding/train_embedding_full_ddp.py +++ b/cookbook/exp/embedding/train_embedding_full_ddp.py @@ -1,14 +1,12 @@ -"""LoRA embedding training with online condenser self-improvement. +"""LoRA embedding training with online compression via frozen vLLM condenser. Architecture (8 GPUs total): - Ranks 0-3 (``model``): Trainable embedding model with LoRA, InfoNCE loss. - - Ranks 4-5 (``condenser_sampler``): Frozen vLLM condenser for online compression. - - Ranks 6-7 (``condenser_model``): Trainable condenser with LoRA for self-improvement. + - Ranks 4-7 (``condenser_sampler``): Frozen vLLM condenser for online compression. -When the condenser sampler truncates (stop_reason='length'), an external OpenAI- -compatible API produces the correct compression. The failure is logged as SFT -training data. A background thread retrains the condenser on accumulated failures -mixed with condense_300K, then syncs weights back to the sampler. +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. Launch: python cookbook/exp/train_embedding_lora_ddp.py @@ -19,6 +17,8 @@ import re import sys import threading +import time +from collections import deque from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any, Dict, List, Literal, Optional @@ -27,7 +27,6 @@ import twinkle from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger -from twinkle.checkpoint_engine import CheckpointEngineManager from twinkle.data_format import SamplingParams from twinkle.dataloader import DataLoader from twinkle.loss import InfonceLoss @@ -55,14 +54,14 @@ # -- GPU placement (8 total) -------------------------------------------------- MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) -CONDENSER_SAMPLER_GPUS = int(os.environ.get('CONDENSER_SAMPLER_GPUS', 2)) -CONDENSER_MODEL_GPUS = int(os.environ.get('CONDENSER_MODEL_GPUS', 2)) -NUM_GPUS = MODEL_GPUS + CONDENSER_SAMPLER_GPUS + CONDENSER_MODEL_GPUS +CONDENSER_SAMPLER_GPUS = int(os.environ.get('CONDENSER_SAMPLER_GPUS', 4)) +NUM_GPUS = MODEL_GPUS + CONDENSER_SAMPLER_GPUS # -- Embedding training hyper-params ------------------------------------------ EMB_MAX_LENGTH = 8192 HARD_NEGATIVES = None -TEMPERATURE = 0.03 +# 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)) LEARNING_RATE = 1.5e-6 @@ -92,16 +91,15 @@ COMPRESS_TOP_P = 0.5 COMPRESS_MAX_MODEL_LEN = 32768 +# Prefetch depth: >1 lets next batch's vLLM/API run while current batch trains. +PREFETCH_WORKERS = int(os.environ.get('PREFETCH_WORKERS', 2)) + # -- 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') - -# -- Condenser retraining knobs ----------------------------------------------- -CONDENSER_DATASET_ID = 'ms://twinkle-kit/condense_300K' -CONDENSER_RETRAIN_SAMPLES = 128 -CONDENSER_RETRAIN_EPOCHS = 3 -CONDENSER_RETRAIN_LR = 1e-5 +# Minimum gap between API calls (seconds); bounds dashscope qps under provider limits. +API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) # -- Output paths ------------------------------------------------------------- OUTPUT_DIR = f'./output/embedding_lora_{BACKEND}' @@ -209,6 +207,20 @@ _sample_counter = 0 _sample_counter_lock = threading.Lock() +# Serializes vLLM sample() across PREFETCH_WORKERS threads. +_sampler_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 @@ -319,11 +331,40 @@ def save_checkpoint(model, name: str): # Compression prompt building # ============================================================================= +# 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 = ( +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.') @@ -347,22 +388,25 @@ def _extract_query_cot(row: Dict[str, Any]): 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) where: + 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. - # At worst-case 1.5 chars/token (CJK), that's ~33k chars. Use 60k as safe English-mix cap. - _MAX_COT_CHARS = 60_000 + # 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: @@ -371,26 +415,32 @@ def _build_compress_prompts(rows: List[Dict[str, Any]]) -> tuple: 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=EMBED_QUERY_Q, text=query) + 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(EMBED_QUERY_Q) - user = COMPRESS_USER.format(query=EMBED_QUERY_COT, text=cot) + 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(EMBED_QUERY_COT) + prompt_queries.append(c_hint) passthrough.append(None) - return prompts, valid_indices, raw_pairs, prompt_queries, passthrough + 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]]: @@ -415,13 +465,24 @@ def _get_first_feature(decoded_text: str, template: Template, role: str) -> Opti # OpenAI API fallback # ============================================================================= -def _is_truncated_compression(text: str) -> bool: - """Detect structurally incomplete output that vLLM may report as stop_reason='stop'. +_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. - The condenser sometimes emits a chat-template token mid-skeleton (which we then - strip), so the visible text ends mid-sentence even though stop_reason!='length'. - The COMPRESS_SYSTEM skeleton mandates a `## More` section ending in a bullet list; - its absence is an unambiguous truncation signal. + 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 @@ -433,11 +494,18 @@ def _is_truncated_compression(text: str) -> bool: 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) @@ -456,61 +524,12 @@ def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[ return content -# ============================================================================= -# Condenser Retrainer (background thread) -# ============================================================================= - -class CondenserRetrainer: - """Async condenser self-improvement: retrains from failures, syncs to sampler.""" - - def __init__(self, condenser_model, ckpt_manager: CheckpointEngineManager, - condenser_sampler): - self._model = condenser_model - self._ckpt_manager = ckpt_manager - self._sampler = condenser_sampler - self._signal = threading.Event() - self._stop = threading.Event() - self._thread = threading.Thread(target=self._loop, daemon=True) - self._condense_300k_cache = None - self._retrain_count = 0 - # Prevents sample() and sync_weights() from running concurrently - self.sampler_lock = threading.Lock() - - def start(self): - self._thread.start() - - def stop(self): - self._stop.set() - self._signal.set() - self._thread.join(timeout=10) - - def notify_failure(self): - self._signal.set() - - def _loop(self): - while not self._stop.is_set(): - self._signal.wait(timeout=60) - if self._stop.is_set(): - break - if not self._signal.is_set(): - continue - self._signal.clear() - try: - self._retrain_and_sync() - except Exception as exc: - logger.error(f'[condenser_retrain] crashed: {exc}') - - def _retrain_and_sync(self): - # Retrain + sync temporarily disabled; failures.jsonl is written directly by _log_failure. - pass - - # ============================================================================= # Main training # ============================================================================= def train(): - # -------- Device groups (3 groups) ---------------------------------------- + # -------- Device groups (2 groups) ---------------------------------------- device_groups = [ DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), @@ -518,15 +537,10 @@ def train(): DeviceGroup(name='condenser_sampler', ranks=list(range(MODEL_GPUS, MODEL_GPUS + CONDENSER_SAMPLER_GPUS)), device_type='GPU'), - DeviceGroup(name='condenser_model', - ranks=list(range(MODEL_GPUS + CONDENSER_SAMPLER_GPUS, NUM_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) - condenser_model_mesh = DeviceMesh.from_sizes( - world_size=CONDENSER_MODEL_GPUS, dp_size=1, fsdp_size=CONDENSER_MODEL_GPUS) twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups) @@ -557,7 +571,7 @@ def train(): setup_optimizer(model, optimizer_steps) model.add_metric(EmbeddingMetric, is_training=True) - # -------- Condenser sampler (2 GPU, vLLM) -------------------------------- + # -------- Condenser sampler (4 GPU, vLLM) -------------------------------- emb_template = Template(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 = Template(model_id=CONDENSE_MODEL_ID, max_length=DATASET_MAX_TOKENS, @@ -582,24 +596,6 @@ def train(): num_samples=1, ) - # -------- Condenser model (2 GPU, trainable full-param) ------------------- - condenser_model = TransformersModel( - model_id=CONDENSE_MODEL_ID, - device_mesh=condenser_model_mesh, - remote_group='condenser_model', - ) - condenser_model.set_optimizer(optimizer_cls='AdamW', lr=CONDENSER_RETRAIN_LR) - - # -------- CheckpointEngineManager: condenser_model → condenser_sampler --- - condenser_ckpt_manager = CheckpointEngineManager( - model=condenser_model, sampler=condenser_sampler) - condenser_ckpt_manager.sync_weights() - - # -------- Background retrainer ------------------------------------------- - retrainer = CondenserRetrainer(condenser_model, condenser_ckpt_manager, - condenser_sampler) - retrainer.start() - # -------- OpenAI API client for fallback --------------------------------- api_client = OpenAIClient( model=COMPRESS_MODEL, @@ -629,7 +625,7 @@ def train(): # -------- Train loop ----------------------------------------------------- def _sample_batch(raw_batch): """Compress via vLLM sampler; fall back to API on truncation.""" - compress_prompts, valid_indices, raw_pairs, prompt_queries, passthrough = \ + compress_prompts, valid_indices, raw_pairs, prompt_queries, passthrough, schemas = \ _build_compress_prompts(raw_batch) if len(compress_prompts) < 4: return None @@ -639,7 +635,7 @@ def _sample_batch(raw_batch): sampler_pos = [ri for ri, p in enumerate(compress_prompts) if p is not None] if sampler_input: try: - with retrainer.sampler_lock: + with _sampler_lock: sampler_responses = condenser_sampler.sample(sampler_input, compress_params) except Exception as exc: logger.warning(f'[sampler] encode overflow in batch, falling back to API: {exc}') @@ -668,7 +664,7 @@ def _sample_batch(raw_batch): # Premature-EOS: model emits chat-template token mid-skeleton, vLLM reports # stop_reason='stop' but the stripped text is structurally incomplete. needs_fallback = (not seq or seq.stop_reason == 'length' - or _is_truncated_compression(text)) + or _is_truncated_compression(text, schemas[ri])) if not needs_fallback: decoded_texts.append(text) continue @@ -676,14 +672,13 @@ def _sample_batch(raw_batch): api_result = _api_compress(api_client, compress_prompts[ri]) # Skip logging when the API itself produced truncated output: an incomplete # gold answer would teach the condenser to imitate broken outputs. - if api_result and not _is_truncated_compression(api_result): + if api_result and not _is_truncated_compression(api_result, schemas[ri]): decoded_texts.append(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]) - retrainer.notify_failure() else: decoded_texts.append('') @@ -711,7 +706,7 @@ def _sample_batch(raw_batch): _start_epoch = cur_step // _batches_per_epoch if cur_step > 0 else 0 _skip_batches_in_epoch = cur_step - _start_epoch * _batches_per_epoch if cur_step > 0 else 0 - prefetch_executor = ThreadPoolExecutor(max_workers=1) + prefetch_executor = ThreadPoolExecutor(max_workers=PREFETCH_WORKERS) for epoch in range(_start_epoch, NUM_EPOCHS): # Skip consumed samples for the resume epoch (shuffle order won't match # exactly, but the correct number of samples is skipped). @@ -720,14 +715,16 @@ def _sample_batch(raw_batch): batch_iter = iter(dataloader) # Reset skip after first resumed epoch _skip_batches_in_epoch = 0 - prefetch_future = None - first_batch = next(batch_iter, None) - if first_batch is not None: - prefetch_future = prefetch_executor.submit(_sample_batch, first_batch) + prefetch_pool: deque = deque() + for _ in range(PREFETCH_WORKERS): + nb = next(batch_iter, None) + if nb is None: + break + prefetch_pool.append(prefetch_executor.submit(_sample_batch, nb)) for raw_batch in batch_iter: - emb_features = prefetch_future.result() if prefetch_future else None - prefetch_future = prefetch_executor.submit(_sample_batch, raw_batch) + emb_features = prefetch_pool.popleft().result() if prefetch_pool else None + prefetch_pool.append(prefetch_executor.submit(_sample_batch, raw_batch)) if emb_features is None: continue @@ -753,16 +750,15 @@ def _sample_batch(raw_batch): if cur_step % SAVE_INTERVAL == 0: save_checkpoint(model, f'step_{cur_step}') - # # Drain last prefetched batch - # if prefetch_future is not None: - # emb_features = prefetch_future.result() + # # Drain remaining prefetched batches + # while prefetch_pool: + # emb_features = prefetch_pool.popleft().result() # if emb_features is not None: # model.forward_backward(inputs=emb_features, task='embedding') # model.clip_grad_and_step() # cur_step += 1 prefetch_executor.shutdown(wait=False) - retrainer.stop() save_checkpoint(model, 'last-checkpoint') From 12471dc85d8f6a6f6392e944d01b3ce0826ed532 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 16 Jun 2026 22:03:33 +0800 Subject: [PATCH 03/26] fix --- .../exp/embedding/build_thinking_rag_index.py | 107 +++++++- .../exp/embedding/train_embedding_full_ddp.py | 235 +++++++++++------- src/twinkle/infra/__init__.py | 8 +- src/twinkle/infra/_ray/ray_helper.py | 10 +- src/twinkle/processor/base.py | 3 +- 5 files changed, 264 insertions(+), 99 deletions(-) diff --git a/cookbook/exp/embedding/build_thinking_rag_index.py b/cookbook/exp/embedding/build_thinking_rag_index.py index 0299aa48..d228a597 100644 --- a/cookbook/exp/embedding/build_thinking_rag_index.py +++ b/cookbook/exp/embedding/build_thinking_rag_index.py @@ -48,15 +48,95 @@ from tqdm import tqdm # --------------------------------------------------------------------------- -# Reuse condenser prompts (single source of truth) +# Compress prompts — MUST match train_embedding_full_ddp.py exactly. # --------------------------------------------------------------------------- _HERE = Path(__file__).resolve().parent -sys.path.insert(0, str(_HERE.parent / 'condenser')) sys.path.insert(0, str(_HERE)) -from make_condenser_dataset import ( # noqa: E402 (after sys.path tweak) - COMPRESS_SYSTEM, - COMPRESS_USER, -) + +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}') + # Default dataset loader is the index-time corpus (broader retrieval profile); # pass --dataset-module dataset_think to fall back to the training mix. from dataset_index import get_dataset as _default_get_dataset # noqa: E402 @@ -83,7 +163,7 @@ EMBED_MODEL_ID = os.environ.get( 'EMBED_MODEL_ID', - '/mnt/workspace/yzhao/tastelikefeet/Qwen3.5-4B-QA-emb-v2', + 'output/embedding_lora_transformers/step_4000', ) CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') @@ -531,8 +611,17 @@ def _flush(rows: List[Dict[str, Any]]) -> None: if not rows: return # Phase 1 — compress query (RAG_QUERY_HINT) and cot (RAG_THINKING_HINT). - q_compressed = _resolve_compressed( - sampler, api, [r['query_raw'] for r in rows], RAG_QUERY_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) kept_rows: List[Dict[str, Any]] = [] diff --git a/cookbook/exp/embedding/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py index 9ac925b6..4e0cd9f2 100644 --- a/cookbook/exp/embedding/train_embedding_full_ddp.py +++ b/cookbook/exp/embedding/train_embedding_full_ddp.py @@ -18,7 +18,6 @@ import sys import threading import time -from collections import deque from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any, Dict, List, Literal, Optional @@ -34,7 +33,7 @@ from twinkle.model import TransformersModel from twinkle.processor import InputProcessor from twinkle.sampler import vLLMSampler -from twinkle.template import Template +from twinkle.template import Qwen3_5Template, Template from twinkle.utils.parallel import PosixFileLock from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient @@ -91,8 +90,8 @@ COMPRESS_TOP_P = 0.5 COMPRESS_MAX_MODEL_LEN = 32768 -# Prefetch depth: >1 lets next batch's vLLM/API run while current batch trains. -PREFETCH_WORKERS = int(os.environ.get('PREFETCH_WORKERS', 2)) +# 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', '') @@ -100,6 +99,9 @@ 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}' @@ -207,9 +209,6 @@ _sample_counter = 0 _sample_counter_lock = threading.Lock() -# Serializes vLLM sample() across PREFETCH_WORKERS threads. -_sampler_lock = threading.Lock() - _api_throttle_lock = threading.Lock() _api_last_call = [0.0] @@ -559,8 +558,9 @@ def train(): 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)}') - dataloader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, shuffle=True) - total_forward_steps = len(dataloader) * NUM_EPOCHS + _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) ---------------------------------------- @@ -572,9 +572,9 @@ def train(): model.add_metric(EmbeddingMetric, is_training=True) # -------- Condenser sampler (4 GPU, vLLM) -------------------------------- - emb_template = Template(model_id=MODEL_ID, max_length=EMB_MAX_LENGTH, enable_thinking=False) + 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 = Template(model_id=CONDENSE_MODEL_ID, max_length=DATASET_MAX_TOKENS, + 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( @@ -596,6 +596,33 @@ def train(): 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, @@ -625,8 +652,10 @@ def train(): # -------- 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 @@ -635,22 +664,29 @@ def _sample_batch(raw_batch): sampler_pos = [ri for ri, p in enumerate(compress_prompts) if p is not None] if sampler_input: try: - with _sampler_lock: - sampler_responses = condenser_sampler.sample(sampler_input, compress_params) + sampler_responses = condenser_sampler.sample(sampler_input, compress_params) except Exception as exc: - logger.warning(f'[sampler] encode overflow in batch, falling back to API: {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[str] = [] + 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.append(passthrough[ri]) + decoded_texts[ri] = passthrough[ri] continue resp = responses[ri] seq = resp.sequences[0] if resp and resp.sequences else None @@ -661,26 +697,33 @@ def _sample_batch(raw_batch): text = text.replace(tok, '') text = text.rstrip() - # Premature-EOS: model emits chat-template token mid-skeleton, vLLM reports - # stop_reason='stop' but the stripped text is structurally incomplete. needs_fallback = (not seq or seq.stop_reason == 'length' or _is_truncated_compression(text, schemas[ri])) if not needs_fallback: - decoded_texts.append(text) - continue - - api_result = _api_compress(api_client, compress_prompts[ri]) - # Skip logging when the API itself produced truncated output: an incomplete - # gold answer would teach the condenser to imitate broken outputs. - if api_result and not _is_truncated_compression(api_result, schemas[ri]): - decoded_texts.append(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]) + decoded_texts[ri] = text else: - decoded_texts.append('') + 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]] = [] @@ -695,68 +738,94 @@ def _sample_batch(raw_batch): if feat_q and feat_c: emb_features.append(feat_q) emb_features.append(feat_c) + _t_feat = time.monotonic() - if len(emb_features) < 4: - return None - return emb_features + 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 - # Compute which epoch and how many batches to skip within that epoch _batches_per_epoch = len(dataloader) - _start_epoch = cur_step // _batches_per_epoch if cur_step > 0 else 0 - _skip_batches_in_epoch = cur_step - _start_epoch * _batches_per_epoch if cur_step > 0 else 0 + _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) - prefetch_executor = ThreadPoolExecutor(max_workers=PREFETCH_WORKERS) + _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): - # Skip consumed samples for the resume epoch (shuffle order won't match - # exactly, but the correct number of samples is skipped). if _skip_batches_in_epoch > 0: - dataloader.skip_consumed_samples(_skip_batches_in_epoch * BATCH_SIZE) + dataloader.skip_consumed_samples(_skip_batches_in_epoch * _mega_batch_size) batch_iter = iter(dataloader) - # Reset skip after first resumed epoch _skip_batches_in_epoch = 0 - prefetch_pool: deque = deque() - for _ in range(PREFETCH_WORKERS): - nb = next(batch_iter, None) - if nb is None: - break - prefetch_pool.append(prefetch_executor.submit(_sample_batch, nb)) - - for raw_batch in batch_iter: - emb_features = prefetch_pool.popleft().result() if prefetch_pool else None - prefetch_pool.append(prefetch_executor.submit(_sample_batch, raw_batch)) - - if emb_features is None: + + first = next(batch_iter, None) + future = prefetch_executor.submit(_sample_batch, first) if first else None + + for raw_mega_batch in batch_iter: + t0 = time.monotonic() + minibatches = future.result() if future else None + t_prefetch = time.monotonic() - t0 + future = prefetch_executor.submit(_sample_batch, raw_mega_batch) + + if not minibatches: continue - model.forward_backward(inputs=emb_features, task='embedding') - model.clip_grad_and_step(gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - 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_forward_steps}, metric: {metric}') - 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 - swanlab.log(log_dict, step=cur_step) - if cur_step % SAVE_INTERVAL == 0: - save_checkpoint(model, f'step_{cur_step}') - - # # Drain remaining prefetched batches - # while prefetch_pool: - # emb_features = prefetch_pool.popleft().result() - # if emb_features is not None: - # model.forward_backward(inputs=emb_features, task='embedding') - # model.clip_grad_and_step() - # cur_step += 1 + 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) save_checkpoint(model, 'last-checkpoint') diff --git a/src/twinkle/infra/__init__.py b/src/twinkle/infra/__init__.py index 83e10d13..d5010138 100644 --- a/src/twinkle/infra/__init__.py +++ b/src/twinkle/infra/__init__.py @@ -700,7 +700,8 @@ def remote_function(dispatch: Union[Literal['slice', 'all', 'slice_dp'], Callabl execute: Literal['first', 'peer', 'all'] = 'all', collect: Union[Literal['none', 'flatten', 'mean', 'sum', 'first', 'last_pp'], Callable] = 'none', sync: bool = False, - lazy_collect: Optional[bool] = None): + lazy_collect: Optional[bool] = None, + timeout: Optional[float] = None): """Patch each method called from remote(which class should be decorated with `remote_class`) with this decorator. Args: @@ -726,6 +727,7 @@ def remote_function(dispatch: Union[Literal['slice', 'all', 'slice_dp'], Callabl sync: If True, use synchronous execution (execute_all_sync) instead of async. Required for methods with NCCL collective operations (e.g., Megatron forward_backward). lazy_collect: Do lazy collect, this boolean value decides whether this function needs lazy collect. If setting to None, it will follow the global setting. + timeout: Timeout in seconds for ray.get() when collecting results. Instance attribute ``_ray_get_timeout`` overrides this. """ # noqa def decorator(func: Callable[..., T1]) -> Callable[..., T1]: @@ -773,7 +775,9 @@ def wrapper(self, *args, **kwargs) -> T1: result = execute_method(func.__name__, _workers_and_args) # This is a result future, call it to get the actual result - result_func = RayHelper.do_get_and_collect_func(_collect_func, collect, result, device_mesh) + _rgt = getattr(self, '_ray_get_timeout', None) or timeout + result_func = RayHelper.do_get_and_collect_func( + _collect_func, collect, result, device_mesh, timeout=_rgt) _local_lazy_collect = _lazy_collect if func.__name__ == '__iter__': # return self diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index 0d8908a3..5cd792c3 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -161,18 +161,20 @@ def get_node_address(): return ip, port @staticmethod - def do_get_and_collect_func(collect_func: Callable, method: Union[str, Callable], futures, device_mesh): + def do_get_and_collect_func(collect_func: Callable, method: Union[str, Callable], futures, device_mesh, + timeout=None): """Return a callable to collect results in the workers.""" class LazyCollect: - def __init__(self, futures, method, collect_func, device_mesh): + def __init__(self, futures, method, collect_func, device_mesh, timeout=None): self._futures = futures self._method = method self._collect_func = collect_func self._is_lazy_collect = True self.device_mesh = device_mesh self._result = None # Cache collected results + self._timeout = timeout def _get_result(self): """Internal method to lazily collect and cache results""" @@ -181,7 +183,7 @@ def _get_result(self): result = [] for future in self._futures: if isinstance(future, ray.ObjectRef): - result.append(ray.get(future)) + result.append(ray.get(future, timeout=self._timeout)) else: result.append(future) self._result = self._collect_func(self._method, result, device_mesh=self.device_mesh) @@ -199,7 +201,7 @@ def __len__(self): """Support len() function""" return len(self._get_result()) - return LazyCollect(futures, method, collect_func, device_mesh) + return LazyCollect(futures, method, collect_func, device_mesh, timeout=timeout) @staticmethod def do_get_and_collect(args, kwargs): diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 8709d98a..750ca69a 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -694,7 +694,8 @@ def is_mm_position_ids(position_ids): result[key] = self._create_4d_attention_mask(values) elif key == 'position_ids' and is_mm_position_ids(values[0]): result[key] = InputProcessor._pad_sequence(values, self.padding_map[key], self.padding_side) - result[key] = result[key].reshape(values[0].shape[0], len(values), -1) + num_axes = values[0].shape[0] + result[key] = result[key].reshape(len(values), num_axes, -1).permute(1, 0, 2).contiguous() elif isinstance(values[0], torch.Tensor): result[key] = InputProcessor._pad_sequence(values, self.padding_map[key], self.padding_side) if result[key].dim() == 1: From e23c2124e4a655b92e6f45f9976dea856b24ce32 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 16 Jun 2026 22:08:13 +0800 Subject: [PATCH 04/26] fix --- cookbook/exp/embedding/train_embedding_full_ddp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cookbook/exp/embedding/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py index 4e0cd9f2..3bb546d7 100644 --- a/cookbook/exp/embedding/train_embedding_full_ddp.py +++ b/cookbook/exp/embedding/train_embedding_full_ddp.py @@ -48,7 +48,7 @@ # 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://twinkle-kit/Qwen3.5-4B-QA-emb') +MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') TEMPLATE_NAME = 'Qwen3_5Template' # -- GPU placement (8 total) -------------------------------------------------- @@ -71,7 +71,7 @@ TOTAL_SAMPLES: Optional[int] = None # Post-build caps on each loader (None = no cap). Applied via .select() before mix. -THINK_CAP: Optional[int] = 200_000 +THINK_CAP: Optional[int] = 400_000 INDEX_CAP: Optional[int] = 400_000 MIX_SHUFFLE_SEED = 42 From d6d4896cd15669902bec654bdbdb7c8100ad1a82 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 17 Jun 2026 09:47:38 +0800 Subject: [PATCH 05/26] fix --- cookbook/exp/embedding/train_embedding_full_ddp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cookbook/exp/embedding/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py index 3bb546d7..5db9f786 100644 --- a/cookbook/exp/embedding/train_embedding_full_ddp.py +++ b/cookbook/exp/embedding/train_embedding_full_ddp.py @@ -63,10 +63,10 @@ TEMPERATURE = 0.07 BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 32)) -LEARNING_RATE = 1.5e-6 +LEARNING_RATE = 1e-5 GRADIENT_ACCUMULATION_STEPS = 1 LOG_INTERVAL = 2 -SAVE_INTERVAL = 4000 +SAVE_INTERVAL = 2000 NUM_EPOCHS = 1 TOTAL_SAMPLES: Optional[int] = None From a498ca289b1a769d1a4bbe3266c83bec041216a6 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 17 Jun 2026 16:33:56 +0800 Subject: [PATCH 06/26] fix --- .../exp/embedding/build_thinking_rag_index.py | 323 ++++++++++++--- cookbook/sample/emb_sample.py | 2 +- cookbook/sample/rag_recall_sample.py | 379 ++++++++++++++++++ 3 files changed, 649 insertions(+), 55 deletions(-) create mode 100644 cookbook/sample/rag_recall_sample.py diff --git a/cookbook/exp/embedding/build_thinking_rag_index.py b/cookbook/exp/embedding/build_thinking_rag_index.py index d228a597..0cdfbe11 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_lora_transformers/step_8000', ) CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') @@ -185,6 +187,16 @@ 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 @@ -419,23 +431,117 @@ 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 [] + + # Build prompts per-item (each text gets its own hint as the query parameter). + prompts = [{'messages': _build_compress_messages(t, h)} for t, h in zip(texts, hints)] + 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) + + results: List[Optional[str]] = [None] * len(texts) + fallback_indices: List[int] = [] + for i, resp in enumerate(responses): + seq = resp.sequences[0] if resp and resp.sequences else None + if seq is None: + fallback_indices.append(i) 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[i] = text else: - results.append(api_text) + fallback_indices.append(i) + + # 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 +651,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 +708,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 +769,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 +787,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 +811,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 +873,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 +1058,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, + p.add_argument('--max-rows', type=int, default=5000, 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 +1095,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/sample/emb_sample.py b/cookbook/sample/emb_sample.py index 6d7e4c59..fc39371d 100644 --- a/cookbook/sample/emb_sample.py +++ b/cookbook/sample/emb_sample.py @@ -32,7 +32,7 @@ # -- Config ------------------------------------------------------------------- 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', 'ms://twinkle-kit/Qwen3.5-4B-QA-emb') +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 diff --git a/cookbook/sample/rag_recall_sample.py b/cookbook/sample/rag_recall_sample.py new file mode 100644 index 00000000..61eccceb --- /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/step_8000') +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() From 4a8f2e88b0035f30b7f322aa92f4decf4d4d89bd Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Thu, 18 Jun 2026 00:29:33 +0800 Subject: [PATCH 07/26] fix --- .gitignore | 2 +- cookbook/exp/embedding/dataset_hard.py | 202 +++++ .../exp/embedding/make_embedding_dataset.py | 735 +++++++++++++++++ .../exp/embedding/train_embedding_full_ddp.py | 750 +++--------------- src/twinkle/template/base.py | 5 +- 5 files changed, 1042 insertions(+), 652 deletions(-) create mode 100644 cookbook/exp/embedding/dataset_hard.py create mode 100644 cookbook/exp/embedding/make_embedding_dataset.py 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/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/make_embedding_dataset.py b/cookbook/exp/embedding/make_embedding_dataset.py new file mode 100644 index 00000000..3417b830 --- /dev/null +++ b/cookbook/exp/embedding/make_embedding_dataset.py @@ -0,0 +1,735 @@ +"""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 hashlib +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', 8)) +SAMPLER_TIMEOUT = float(os.environ.get('SAMPLER_TIMEOUT', 300)) + +# -- Output ------------------------------------------------------------------- +OUTPUT_DIR = os.environ.get('EMB_DATASET_OUTPUT', './output/embedding_dataset') +FAILURE_LOG = os.environ.get('FAILURE_LOG', f'{OUTPUT_DIR}/failures.jsonl') +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_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(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 + + +_failure_lock: Optional[PosixFileLock] = None + + +def _log_failure(source_text: str, query: str, compressed: str, row_id: str): + 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'{row_id}__{qhash}', + 'source': 'compress_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) + + +# ============================================================================= +# 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 + + # API fallback for failed items + 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 + + 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'] + 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 + results.append({ + 'anchor_text': q_text, + 'positive_text': c_text, + 'negative_texts': [], + 'source': meta[pair_idx]['source'], + }) + 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'] + results = [] + offset = 0 + for gs in group_sizes: + q_text = decoded[offset] + c_text = decoded[offset + 1] + if not q_text or not c_text: + offset += gs + continue + neg_texts = [] + for ni in range(2, gs): + nt = decoded[offset + ni] + if nt: + neg_texts.append(nt) + results.append({ + 'anchor_text': q_text, + 'positive_text': c_text, + 'negative_texts': neg_texts, + 'source': source_type, + }) + 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] = [] + + for row in rows: + query, cot = _extract_query_cot(row) + if not query or not cot or len(cot) > _MAX_COT_CHARS: + continue + neg_json = row.get('negatives_json', '') + negatives = json.loads(neg_json) if neg_json else [] + 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', '')) + + if not prompts: + return {'hard': True, 'prompts': [], 'group_sizes': [], 'row_ids': [], + 'decoded': [], 'fallback_indices': [], 'source_type': source_type} + + 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} + + +# ============================================================================= +# 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 + # Ensure negatives_json column exists + if 'negatives_json' not in ds_think.column_names: + ds_think = ds_think.add_column('negatives_json', [''] * len(ds_think)) + 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)) + if 'negatives_json' not in ds_index.column_names: + ds_index = ds_index.add_column('negatives_json', [''] * len(ds_index)) + 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: + for i in range(n_hard): + row = ds_hard_raw[i] + hard_rows_list.append({ + 'id': row['id'], + 'messages': [ + {'role': 'user', 'content': row['query']}, + {'role': 'assistant', 'reasoning_content': row['cot'], + 'content': row.get('response', '') or ''}, + ], + 'negatives_json': json.dumps(row['negatives'], ensure_ascii=False), + }) + + # Batch-convert HF Datasets to list-of-dicts + def _ds_to_rows(ds): + cols = {k: ds[k] for k in ds.column_names} + return [{k: cols[k][i] for k in cols} for i in range(len(ds))] + + 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 + + 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 + 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() + + 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 in f: + if line.strip(): + all_results.append(json.loads(line)) + 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], + }) + 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 5db9f786..75f86e86 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) -------------------------------------------------- +# -- GPU placement ------------------------------------------------------------ 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 # -- 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', 16)) 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', './output/embedding_dataset/dataset') 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,29 @@ 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 datasets import Dataset as HFDataset + logger.info(f'[data] loading pre-compressed dataset from {DATASET_PATH}') + dataset = HFDataset.load_from_disk(DATASET_PATH) + dataset = dataset.shuffle(seed=MIX_SHUFFLE_SEED) + logger.info(f'[data] {len(dataset)} rows loaded') + + # -- Compute steps -------------------------------------------------------- + _target = BATCH_SIZE * 2 # features per minibatch (anchor + positive pairs) + # Estimate: each row produces ~2 features (+ negatives); conservative estimate + 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,234 +194,73 @@ 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 + # Split into minibatches at group boundaries + group_starts = [i for i, f in enumerate(features) if f.get('labels') == [1]] + minibatches = [] + mb_start_idx = 0 + for gi in range(len(group_starts)): + pos = group_starts[gi] + if pos - group_starts[mb_start_idx] >= _target: + minibatches.append(features[group_starts[mb_start_idx]:pos]) + mb_start_idx = gi + if mb_start_idx < len(group_starts): + minibatches.append(features[group_starts[mb_start_idx]:]) + minibatches = [mb for mb in minibatches if len(mb) >= 4] + 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) + 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}') + 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: @@ -808,27 +270,15 @@ def _sample_batch(raw_batch): except (ValueError, TypeError): pass log_dict['epoch'] = epoch - log_dict['prefetch_sec'] = round(t_prefetch, 2) - log_dict['train_sec'] = round(t_train, 2) + 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}') - 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) + t_encode = 0.0 + save_checkpoint(model, 'last-checkpoint') + logger.info(f'Training complete. Final step: {cur_step}') if __name__ == '__main__': diff --git a/src/twinkle/template/base.py b/src/twinkle/template/base.py index 26c2e4f2..bf981f29 100644 --- a/src/twinkle/template/base.py +++ b/src/twinkle/template/base.py @@ -658,7 +658,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]): From dd10f2c70d467227c11ce91dfa555dffee2f1607 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Thu, 18 Jun 2026 14:35:34 +0800 Subject: [PATCH 08/26] fix --- .../exp/embedding/make_embedding_dataset.py | 141 ++++++++++-------- 1 file changed, 82 insertions(+), 59 deletions(-) diff --git a/cookbook/exp/embedding/make_embedding_dataset.py b/cookbook/exp/embedding/make_embedding_dataset.py index 3417b830..1aaa942f 100644 --- a/cookbook/exp/embedding/make_embedding_dataset.py +++ b/cookbook/exp/embedding/make_embedding_dataset.py @@ -8,7 +8,6 @@ Launch (8 GPUs — 4 for vLLM condenser): python cookbook/exp/embedding/make_embedding_dataset.py """ -import hashlib import json import os import re @@ -61,12 +60,11 @@ 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', 8)) +API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 32)) SAMPLER_TIMEOUT = float(os.environ.get('SAMPLER_TIMEOUT', 300)) # -- Output ------------------------------------------------------------------- OUTPUT_DIR = os.environ.get('EMB_DATASET_OUTPUT', './output/embedding_dataset') -FAILURE_LOG = os.environ.get('FAILURE_LOG', f'{OUTPUT_DIR}/failures.jsonl') RESULTS_JSONL = f'{OUTPUT_DIR}/results.jsonl' PROGRESS_FILE = f'{OUTPUT_DIR}/progress.json' @@ -243,16 +241,30 @@ def _is_truncated_compression(text: str, schema: str = 'new') -> bool: return False -_api_throttle_lock = threading.Lock() -_api_last_call = [0.0] +_api_semaphore = threading.Semaphore(API_CONCURRENCY) +_api_bucket_lock = threading.Lock() +_api_tokens = [float(API_CONCURRENCY)] +_api_last_refill = [time.monotonic()] 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() + """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]: @@ -273,33 +285,6 @@ def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[ return content -_failure_lock: Optional[PosixFileLock] = None - - -def _log_failure(source_text: str, query: str, compressed: str, row_id: str): - 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'{row_id}__{qhash}', - 'source': 'compress_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) - - # ============================================================================= # Core compression logic # ============================================================================= @@ -415,7 +400,8 @@ def _compress_batch_phase2( is_hard = state.get('hard', False) meta = state.get('meta') # None for hard - # API fallback for failed items + # Track which prompts used API fallback + api_set: set = set() if fallback_indices: api_futures = {} with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: @@ -427,7 +413,9 @@ def _compress_batch_phase2( 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) @@ -436,17 +424,24 @@ def _compress_batch_phase2( 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 @@ -455,24 +450,38 @@ 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 gs in group_sizes: + 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 @@ -491,13 +500,13 @@ def _compress_hard_phase1( 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 - neg_json = row.get('negatives_json', '') - negatives = json.loads(neg_json) if neg_json else [] + negatives = row.get('negatives') or [] valid_negs = [n for n in negatives if n and len(n) <= _MAX_COT_CHARS] @@ -519,10 +528,12 @@ def _compress_hard_phase1( ]}) 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} + 'decoded': [], 'fallback_indices': [], 'source_type': source_type, + 'raw_groups': []} try: responses = condenser_sampler.sample(prompts, compress_params) @@ -548,7 +559,8 @@ def _compress_hard_phase1( return {'hard': True, 'prompts': prompts, 'group_sizes': group_sizes, 'row_ids': row_ids, 'decoded': decoded, - 'fallback_indices': fallback_indices, 'source_type': source_type} + 'fallback_indices': fallback_indices, 'source_type': source_type, + 'raw_groups': raw_groups} # ============================================================================= @@ -572,17 +584,12 @@ def main(): 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 - # Ensure negatives_json column exists - if 'negatives_json' not in ds_think.column_names: - ds_think = ds_think.add_column('negatives_json', [''] * len(ds_think)) 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)) - if 'negatives_json' not in ds_index.column_names: - ds_index = ds_index.add_column('negatives_json', [''] * len(ds_index)) logger.info(f'[load] index={len(ds_index)}') ds_hard_raw = get_dataset_hard(max_negatives=HARD_MAX_NEGATIVES, load_from_cache_file=True) @@ -594,22 +601,25 @@ def main(): # 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): - row = ds_hard_raw[i] hard_rows_list.append({ - 'id': row['id'], + 'id': h_ids[i], 'messages': [ - {'role': 'user', 'content': row['query']}, - {'role': 'assistant', 'reasoning_content': row['cot'], - 'content': row.get('response', '') or ''}, + {'role': 'user', 'content': h_queries[i]}, + {'role': 'assistant', 'reasoning_content': h_cots[i], + 'content': h_responses[i] or ''}, ], - 'negatives_json': json.dumps(row['negatives'], ensure_ascii=False), + 'negatives': h_negatives[i], }) # Batch-convert HF Datasets to list-of-dicts def _ds_to_rows(ds): - cols = {k: ds[k] for k in ds.column_names} - return [{k: cols[k][i] for k in cols} for i in range(len(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) @@ -663,6 +673,11 @@ def _save_progress(): # -- 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 @@ -678,7 +693,7 @@ def _process_source(rows, source_type, label): pending = None # (future, batch_start, batch_len) def _drain_pending(): - nonlocal total_flushed + nonlocal total_flushed, pending if pending is None: return fut, p_start, p_len = pending @@ -687,6 +702,7 @@ def _drain_pending(): 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] @@ -714,15 +730,22 @@ def _drain_pending(): logger.info(f'[save] converting results.jsonl to HF Dataset...') all_results = [] with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: - for line in f: - if line.strip(): + 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') From 04aaf76d6e85ee5f27316df0e0b5de028d776a8c Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Fri, 19 Jun 2026 10:19:26 +0800 Subject: [PATCH 09/26] fix --- .../exp/embedding/make_embedding_dataset.py | 2 +- .../exp/embedding/train_embedding_full_ddp.py | 83 ++++++++----------- 2 files changed, 34 insertions(+), 51 deletions(-) diff --git a/cookbook/exp/embedding/make_embedding_dataset.py b/cookbook/exp/embedding/make_embedding_dataset.py index 1aaa942f..847f222f 100644 --- a/cookbook/exp/embedding/make_embedding_dataset.py +++ b/cookbook/exp/embedding/make_embedding_dataset.py @@ -60,7 +60,7 @@ 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', 32)) +API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 24)) SAMPLER_TIMEOUT = float(os.environ.get('SAMPLER_TIMEOUT', 300)) # -- Output ------------------------------------------------------------------- diff --git a/cookbook/exp/embedding/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py index 75f86e86..d9b35f6c 100644 --- a/cookbook/exp/embedding/train_embedding_full_ddp.py +++ b/cookbook/exp/embedding/train_embedding_full_ddp.py @@ -31,14 +31,14 @@ MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') # -- GPU placement ------------------------------------------------------------ -MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 8)) # -- Embedding training hyper-params ------------------------------------------ EMB_MAX_LENGTH = 8192 HARD_NEGATIVES = None TEMPERATURE = 0.07 -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 16)) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 64)) LEARNING_RATE = 1e-5 GRADIENT_ACCUMULATION_STEPS = 1 LOG_INTERVAL = 2 @@ -46,7 +46,7 @@ NUM_EPOCHS = 1 # -- Dataset path (output of make_embedding_dataset.py) ----------------------- -DATASET_PATH = os.environ.get('EMB_DATASET_PATH', './output/embedding_dataset/dataset') +DATASET_PATH = os.environ.get('EMB_DATASET_PATH', 'ms://twinkle-kit/qth-embedding') MIX_SHUFFLE_SEED = 42 # -- Resume from checkpoint --------------------------------------------------- @@ -173,15 +173,13 @@ def train(): twinkle.initialize(mode='ray', nproc_per_node=MODEL_GPUS, groups=device_groups) # -- Load pre-compressed dataset ------------------------------------------ - from datasets import Dataset as HFDataset + from twinkle.dataset import Dataset as TwinkleDataset, DatasetMeta logger.info(f'[data] loading pre-compressed dataset from {DATASET_PATH}') - dataset = HFDataset.load_from_disk(DATASET_PATH) - dataset = dataset.shuffle(seed=MIX_SHUFFLE_SEED) + 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 -------------------------------------------------------- - _target = BATCH_SIZE * 2 # features per minibatch (anchor + positive pairs) - # Estimate: each row produces ~2 features (+ negatives); conservative estimate rows_per_step = BATCH_SIZE total_steps = (len(dataset) // rows_per_step) * NUM_EPOCHS optimizer_steps = total_steps // GRADIENT_ACCUMULATION_STEPS @@ -234,48 +232,33 @@ def train(): if len(features) < 4: continue - # Split into minibatches at group boundaries - group_starts = [i for i, f in enumerate(features) if f.get('labels') == [1]] - minibatches = [] - mb_start_idx = 0 - for gi in range(len(group_starts)): - pos = group_starts[gi] - if pos - group_starts[mb_start_idx] >= _target: - minibatches.append(features[group_starts[mb_start_idx]:pos]) - mb_start_idx = gi - if mb_start_idx < len(group_starts): - minibatches.append(features[group_starts[mb_start_idx]:]) - minibatches = [mb for mb in minibatches if len(mb) >= 4] - - 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 - - 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}') - t_encode = 0.0 + 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') logger.info(f'Training complete. Final step: {cur_step}') From 2e3458819a30d3736a4bb58adc48fb33b223093a Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Sun, 28 Jun 2026 10:07:53 +0800 Subject: [PATCH 10/26] wip --- cookbook/exp/embedding/ablation_all.sh | 84 ++ cookbook/exp/embedding/ablation_direct.sh | 12 + .../embedding/ablation_rag_api_condenser.sh | 17 + .../embedding/ablation_rag_local_condenser.sh | 18 + cookbook/exp/embedding/ablation_rag_raw.sh | 15 + .../exp/embedding/build_thinking_rag_index.py | 57 +- cookbook/exp/embedding/eval_gpqa_rag.py | 850 ++++++++++++++++++ cookbook/exp/embedding/eval_rag_recall.py | 187 ++++ .../exp/embedding/train_embedding_full_ddp.py | 2 + cookbook/sample/rag_recall_sample.py | 2 +- 10 files changed, 1219 insertions(+), 25 deletions(-) create mode 100755 cookbook/exp/embedding/ablation_all.sh create mode 100755 cookbook/exp/embedding/ablation_direct.sh create mode 100755 cookbook/exp/embedding/ablation_rag_api_condenser.sh create mode 100755 cookbook/exp/embedding/ablation_rag_local_condenser.sh create mode 100755 cookbook/exp/embedding/ablation_rag_raw.sh create mode 100644 cookbook/exp/embedding/eval_gpqa_rag.py create mode 100644 cookbook/exp/embedding/eval_rag_recall.py 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 0cdfbe11..a71bae06 100644 --- a/cookbook/exp/embedding/build_thinking_rag_index.py +++ b/cookbook/exp/embedding/build_thinking_rag_index.py @@ -165,7 +165,7 @@ EMBED_MODEL_ID = os.environ.get( 'EMBED_MODEL_ID', - 'output/embedding_lora_transformers/step_8000', + 'output/embedding_full_transformers/last-checkpoint', ) CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') @@ -202,25 +202,27 @@ # 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', '') @@ -501,8 +503,14 @@ def _resolve_compressed_multi(sampler: vLLMSampler, api: Optional[OpenAIClient], 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 in zip(texts, hints)] + 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, @@ -511,22 +519,23 @@ def _resolve_compressed_multi(sampler: vLLMSampler, api: Optional[OpenAIClient], ) # Single vLLM batch call — the key throughput win. - responses = sampler.sample(prompts, params) + responses = sampler.sample(prompts, params) if prompts else [] results: List[Optional[str]] = [None] * len(texts) fallback_indices: List[int] = [] - for i, resp in enumerate(responses): + 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(i) + fallback_indices.append(orig_idx) continue 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[i] = text + results[orig_idx] = text else: - fallback_indices.append(i) + fallback_indices.append(orig_idx) # Concurrent API fallback for failed items. if fallback_indices and api is not None: @@ -1064,7 +1073,7 @@ def parse_args() -> argparse.Namespace: '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=5000, + 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=128, diff --git a/cookbook/exp/embedding/eval_gpqa_rag.py b/cookbook/exp/embedding/eval_gpqa_rag.py new file mode 100644 index 00000000..59957d2d --- /dev/null +++ b/cookbook/exp/embedding/eval_gpqa_rag.py @@ -0,0 +1,850 @@ +"""AoPS math competition evaluation: direct vs RAG-augmented with Qwen3.5-4B. + +Two modes: + - ``direct``: The model solves problems directly (4 GPUs, TP=4). + - ``rag``: Retrieve top-k thinking traces from LanceDB as 1-shot + examples, then solve (8 GPUs: DP=4 embedding + TP=4 vLLM). + +Only problems with ``metadata.boxed == True`` are used (auto-gradable via +``\\boxed{...}`` extraction). A random subset is sampled for efficiency. + +Launch examples: + # Direct (4 GPUs, default 500 problems) + python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct + + # RAG-augmented (8 GPUs) + python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag + + # Smaller sample for quick test + python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct --n 100 + + # RAG with condenser (retrieves thinking_raw, compresses with local 4B / API) + EVAL_CONDENSER_GPUS=2 python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag --condense + + # RAG without condenser (uses thinking_raw by default, truncated to max-trace-len) + python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag + + # RAG with pre-compressed field (opt-in, not recommended for reader LM) + python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag --use-cot-compressed +""" +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 + +# --------------------------------------------------------------------------- +# 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') + + +# --------------------------------------------------------------------------- +# 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() + 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 + time.sleep(wait) + finally: + _api_semaphore.release() + + +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 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, +) -> 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] + 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) + 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)}, (B), etc. + 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 = 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 + + +def answers_match(predicted: str, reference: str) -> bool: + """Check if two math answers are equivalent.""" + 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 + return _try_numeric_equal(norm_p, norm_r) + + +# --------------------------------------------------------------------------- +# 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 + + +# --------------------------------------------------------------------------- +# 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{}.' +) + + +def build_direct_prompt(problem: str) -> Dict[str, Any]: + return { + 'messages': [ + {'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}, + ] + } + + +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}) + 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 extra candidates when decontamination is active + fetch_limit = top_k + 50 if decontam_threshold > 0 else top_k + all_examples: List[List[Dict[str, str]]] = [] + decontam_skipped = 0 + for qi, vec in enumerate(query_vecs): + 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 = [] + 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 + # 13-gram decontamination: skip if retrieved query overlaps with problem + if decontam_threshold > 0 and problem_text and q: + ng_sim = _ngram_jaccard(problem_text, q) + if ng_sim > decontam_threshold: + decontam_skipped += 1 + continue + examples.append({'query': q, 'thinking': t, '_sim': round(sim, 4), + '_raw_trace_len': len(t)}) + all_examples.append(examples) + 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='direct') + p.add_argument('--n', type=int, default=500, + help='Number of problems to sample (0 = all boxed).') + p.add_argument('--db-path', default='./output/thinking_rag/lance.db') + p.add_argument('--table', default='thinking_traces') + p.add_argument('--top-k', type=int, default=3) + 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) + p.add_argument('--rag-fallback-sim', type=float, default=0.70, + help='Fallback to direct prompt when best retrieval similarity ' + 'is below this threshold (avoids weak-trace loops).') + 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('--max-trace-len', type=int, default=12000) + p.add_argument('--condense', action='store_true', + help='Enable condenser re-compression on 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('--output', default=None) + args = p.parse_args() + + if args.output is None: + args.output = f'./output/thinking_rag/aops_{args.mode}_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 + + records = load_aops(n=args.n, seed=args.seed) + + is_rag = (args.mode == 'rag') + + condenser_gpus = int(os.environ.get('EVAL_CONDENSER_GPUS', 0)) if args.condense else 0 + + 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') + + correct_count = 0 + total_count = 0 + 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') + + for batch_start in range(0, len(records), args.batch_size): + batch_end = min(batch_start + args.batch_size, len(records)) + batch = records[batch_start:batch_end] + + if is_rag: + 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) + # Strip structural markers from cot_compressed field + if args.use_cot_compressed: + for exs in all_examples: + for ex in exs: + ex['thinking'] = _strip_condenser_markers(ex['thinking']) + + # Condenser re-compression + 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) + + prompts = [] + for r, examples in zip(batch, all_examples): + if not examples: + prompts.append(build_direct_prompt(r['problem'])) + else: + # Drop traces exceeding max_trace_len instead of truncating + filtered = [{'query': ex['query'], 'thinking': ex['thinking']} + for ex in examples + if len(ex['thinking']) <= args.max_trace_len] + # Fallback to direct if best similarity is too low + best_sim = max(ex.get('_sim', 0.0) for ex in examples) + if filtered and best_sim >= args.rag_fallback_sim: + prompts.append(build_rag_prompt(r['problem'], filtered)) + else: + prompts.append(build_direct_prompt(r['problem'])) + else: + 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 + + debug_rec = { + 'idx': batch_start + i, + 'reference_answer': rec['reference_answer'], + 'predicted': predicted, + 'is_correct': is_correct, + 'problem': rec['problem'], + 'model_output': raw_output, + } + if is_rag: + 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', '') + 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' [{batch_end}/{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'AoPS Math — mode={args.mode}, model={GEN_MODEL_ID}') + print(f' n={total_count}, seed={args.seed}') + print(f'{"=" * 60}') + print(f'Overall accuracy: {overall_acc:.4f} ({correct_count}/{total_count})') + + out_f.close() + print(f'\n[output] {len(debug_records)} records saved to {args.output}') + + +if __name__ == '__main__': + main() 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/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py index d9b35f6c..97ab3b12 100644 --- a/cookbook/exp/embedding/train_embedding_full_ddp.py +++ b/cookbook/exp/embedding/train_embedding_full_ddp.py @@ -261,6 +261,8 @@ def train(): 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}') diff --git a/cookbook/sample/rag_recall_sample.py b/cookbook/sample/rag_recall_sample.py index 61eccceb..691a69f6 100644 --- a/cookbook/sample/rag_recall_sample.py +++ b/cookbook/sample/rag_recall_sample.py @@ -37,7 +37,7 @@ # --------------------------------------------------------------------------- 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/step_8000') + '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 From ea932a6fd9cd54ae9d021b91098bb94ce93417a9 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Mon, 29 Jun 2026 11:42:16 +0800 Subject: [PATCH 11/26] wip --- cookbook/exp/rl/rag_hint_grpo.py | 739 +++++++++++++++++++++++++++++++ 1 file changed, 739 insertions(+) create mode 100644 cookbook/exp/rl/rag_hint_grpo.py diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py new file mode 100644 index 00000000..70917d46 --- /dev/null +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -0,0 +1,739 @@ +"""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 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.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', 4)) +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 2)) +NUM_GPUS = CONDENSER_GPUS + EMB_GPUS + SAMPLER_GPUS + MODEL_GPUS + +# Training hyperparams +NUM_GENERATIONS = int(os.environ.get('NUM_GENERATIONS', 16)) +MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 65536)) +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', 4)) +MINI_BATCH_SIZE = int(os.environ.get('MINI_BATCH_SIZE', 4)) +MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 2)) +GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 1)) +SAVE_STEPS = int(os.environ.get('SAVE_STEPS', 100)) + +# 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.65)) +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)) + +# Forced analysis prefix appended at the start of assistant response +ANALYSIS_PREFIX = ( + "Let me think step by step. First, I will analyze the example above: " + "identify which steps and concepts are CORRECT and APPLICABLE to this problem, " + "and explicitly discard any reasoning that is WRONG, IRRELEVANT, or based on " + "assumptions that do not hold here. Then I will solve the problem using only " + "the validated useful parts:\n\n" +) + +# ============================================================================ +# 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. ' + 'Analyze them critically — identify which steps/concepts are applicable ' + 'and which may not apply. Then solve the actual problem step by step. ' + 'Put your final answer inside \\boxed{{}}.\n\n' +) + +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{}.' +) + + +# ============================================================================ +# 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() + + +# ============================================================================ +# Embedding & Retrieval +# ============================================================================ +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_topk(tbl, query_vecs: np.ndarray, sim_threshold: float + ) -> List[List[Dict[str, str]]]: + """Retrieve top-K thinking_raw per query. Returns empty list if none above threshold.""" + results = [] + for vec in query_vecs: + hits = ( + tbl.search(vec.astype(np.float32).tolist()) + .metric('dot') + .limit(TOP_K + 10) + .select(['query_raw', 'thinking_raw', '_distance']) + .to_list() + ) + matched = [] + for h in hits: + sim = 1.0 - h.get('_distance', 0.0) + if sim < sim_threshold: + continue + t = h.get('thinking_raw', '') + if t: + matched.append({'query': h.get('query_raw', ''), 'thinking': t, 'sim': sim}) + if len(matched) >= TOP_K: + break + results.append(matched) + 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 '' + + @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 = 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() + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\\circ|\^\{\\circ\}|°', 'deg', s) + + 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'(? bool: + 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 + 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 + + @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 + return cls._try_numeric_equal(norm_p, norm_r) + + 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 + + +def compute_rewards(trajectories: List[Dict[str, Any]] + ) -> Tuple[List[float], List[float], List[float]]: + acc_fn = AoPSAccuracyReward() + fmt_fn = FormatReward() + acc = acc_fn(trajectories) + fmt = fmt_fn(trajectories) + total = [a + f for a, f in zip(acc, fmt)] + 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', gpus_per_worker=SAMPLER_GPUS), + DeviceGroup(name='model', ranks=list(range(model_start, NUM_GPUS)), + device_type='GPU'), + ] + + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + 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) + 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('GRPOLoss', epsilon=0.2) + model.set_processor(InputProcessor, padding_free=True) + model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=True) + + # -- Rollout sampler -- + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 65536, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=True) + + # -- Embedding model -- + emb_model = TransformersModel( + model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, remote_group='emb_model') + emb_model.set_processor(InputProcessor) + 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 -- + from concurrent.futures import ThreadPoolExecutor + prefetch_pool = ThreadPoolExecutor(max_workers=1) + + 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 = 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, SIM_THRESHOLD) + + # 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]}) + + # Build prompts + rag_prompts = [] + for i, prob in enumerate(problems): + examples = condensed_examples[i] + if examples: + 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'])) + sys_content = ''.join(parts) + else: + sys_content = SYSTEM_DIRECT + + prompt_feature = { + 'messages': [ + {'role': 'system', 'content': sys_content}, + {'role': 'user', 'content': prob}, + ], + 'user_data': [('ground_truth', ground_truths[i])], + 'assistant_prefix': ANALYSIS_PREFIX if examples else '', + } + rag_prompts.append(prompt_feature) + + return rag_prompts + + # Submit first batch prefetch + 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 + + while pending_future is not None: + if optim_step >= MAX_STEPS: + break + + metrics.reset() + rag_prompts = pending_future.result() + + # 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) + + 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) + + 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() + + # ---- Mini-batch training ---- + total_completions = len(all_input_data) + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + mb_inputs = all_input_data[mb_start:mb_end] + mb_old_logps = all_old_logps[mb_start:mb_end] + mb_advantages = advantages[mb_start:mb_end] + + model.forward_backward( + inputs=mb_inputs, + old_logps=mb_old_logps, + advantages=mb_advantages, + micro_batch_size=MICRO_BATCH_SIZE, + ) + 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}') + + 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}') + + prefetch_pool.shutdown(wait=False) + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('rag-hint-grpo-final') + + +if __name__ == '__main__': + main() From 8896c84434882423930ed1cbfa3ef264289d8f7e Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Mon, 29 Jun 2026 11:50:18 +0800 Subject: [PATCH 12/26] fix --- cookbook/exp/rl/rag_hint_grpo.py | 275 +++++++++++++++++++++---------- 1 file changed, 185 insertions(+), 90 deletions(-) diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index 70917d46..73aeb4b7 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -16,6 +16,7 @@ Launch: python cookbook/exp/rl/rag_hint_grpo.py """ +import json import os import re import random @@ -89,6 +90,13 @@ 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 = ( "Let me think step by step. First, I will analyze the example above: " @@ -175,6 +183,30 @@ def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Opti # ============================================================================ # 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}, @@ -202,29 +234,47 @@ def get_embeddings(model: TransformersModel, template: Qwen3_5Template, return emb[:n] if pad_n else emb -def retrieve_topk(tbl, query_vecs: np.ndarray, sim_threshold: float - ) -> List[List[Dict[str, str]]]: - """Retrieve top-K thinking_raw per query. Returns empty list if none above threshold.""" +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 = [] - for vec in query_vecs: + decontam_skipped = 0 + for qi, vec in enumerate(query_vecs): hits = ( tbl.search(vec.astype(np.float32).tolist()) .metric('dot') - .limit(TOP_K + 10) + .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 t: - matched.append({'query': h.get('query_raw', ''), 'thinking': t, 'sim': sim}) - if len(matched) >= TOP_K: - break + 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 @@ -460,7 +510,8 @@ def main(): model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) model.set_loss('GRPOLoss', epsilon=0.2) model.set_processor(InputProcessor, padding_free=True) - model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=True) + model.set_template('Qwen3_5Template', model_id=MODEL_ID, + enable_thinking=True, max_length=65536) # -- Rollout sampler -- sampler = vLLMSampler( @@ -472,7 +523,8 @@ def main(): device_mesh=sampler_mesh, remote_group='sampler', ) - sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=True) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, + enable_thinking=True, max_length=65536) # -- Embedding model -- emb_model = TransformersModel( @@ -536,7 +588,6 @@ def main(): logger.info(get_device_placement()) # -- Prefetch: overlap RAG data preparation with training -- - from concurrent.futures import ThreadPoolExecutor prefetch_pool = ThreadPoolExecutor(max_workers=1) def prepare_rag_batch(batch): @@ -561,7 +612,7 @@ def prepare_rag_batch(batch): # Embed & retrieve query_vecs = get_embeddings(emb_model, emb_template, problems, EMB_GPUS) - retrieved = retrieve_topk(tbl, query_vecs, SIM_THRESHOLD) + retrieved = retrieve_topk(tbl, query_vecs, problems, SIM_THRESHOLD) # Condense (batch local vLLM + API fallback) condensed_examples: List[List[Dict[str, str]]] = [[] for _ in range(len(problems))] @@ -615,11 +666,16 @@ def _fallback(ci): condensed_examples[idx].append( {'query': ret['query'], 'thinking': ret['thinking'][:MAX_TRACE_LEN]}) - # Build prompts + # Build prompts with rag_fallback_sim check rag_prompts = [] + rag_debug_records = [] for i, prob in enumerate(problems): examples = condensed_examples[i] - if examples: + 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: parts = [SYSTEM_WITH_RAG_HEADER] for eidx, ex in enumerate(examples, 1): parts.append(EXAMPLE_TEMPLATE.format( @@ -636,13 +692,33 @@ def _fallback(ci): {'role': 'user', 'content': prob}, ], 'user_data': [('ground_truth', ground_truths[i])], - 'assistant_prefix': ANALYSIS_PREFIX if examples else '', + 'assistant_prefix': ANALYSIS_PREFIX if use_rag else '', } rag_prompts.append(prompt_feature) - return rag_prompts + # Diagnostic record + debug_rec = { + 'problem': prob[:200], + 'ground_truth': ground_truths[i], + 'best_sim': round(best_sim, 4), + 'num_retrieved': len(rets), + 'num_condensed': len(examples), + 'use_rag': use_rag, + } + if rets: + debug_rec['top_retrieved_query'] = rets[0]['query'][:200] + if examples: + debug_rec['condensed_len'] = len(examples[0].get('thinking', '')) + rag_debug_records.append(debug_rec) + + 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, 'a', encoding='utf-8') + logger.info(f'[rag] diagnostics → {rag_log_path}') + batch_iter = iter(dataloader) pending_future = None try: @@ -651,86 +727,105 @@ def _fallback(ci): except StopIteration: pass - while pending_future is not None: - if optim_step >= MAX_STEPS: - break + try: + while pending_future is not None: + if optim_step >= MAX_STEPS: + break - metrics.reset() - rag_prompts = pending_future.result() + metrics.reset() + rag_prompts, rag_debug_records = pending_future.result() - # 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 + # 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() - # ---- Expand for NUM_GENERATIONS and sample ---- - expand_prompts = [] - for prompt in rag_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) - - 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() - - # ---- Mini-batch training ---- - total_completions = len(all_input_data) - for mb_start in range(0, total_completions, MINI_BATCH_SIZE): - mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) - mb_inputs = all_input_data[mb_start:mb_end] - mb_old_logps = all_old_logps[mb_start:mb_end] - mb_advantages = advantages[mb_start:mb_end] - - model.forward_backward( - inputs=mb_inputs, - old_logps=mb_old_logps, - advantages=mb_advantages, - micro_batch_size=MICRO_BATCH_SIZE, + # 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) + + 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) + + # 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') + rag_log_f.flush() + + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={ + 'total': total_rewards, + 'format': format_rewards, + 'accuracy': accuracy_rewards, + }, ) - 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}') + # ---- GRPO advantage ---- + advantages = advantage_fn( + total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + + # ---- Mini-batch training ---- + total_completions = len(all_input_data) + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + mb_inputs = all_input_data[mb_start:mb_end] + mb_old_logps = all_old_logps[mb_start:mb_end] + mb_advantages = advantages[mb_start:mb_end] + + model.forward_backward( + inputs=mb_inputs, + old_logps=mb_old_logps, + advantages=mb_advantages, + micro_batch_size=MICRO_BATCH_SIZE, + ) + 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}') - 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}') + 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() - prefetch_pool.shutdown(wait=False) logger.info(f'Training completed. optim_steps={optim_step}') model.save('rag-hint-grpo-final') From b2442eea3cc52ab6f678efa90decf4cd3cbc8ddb Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Mon, 29 Jun 2026 14:06:35 +0800 Subject: [PATCH 13/26] wip --- cookbook/exp/rl/rag_hint_grpo.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index 73aeb4b7..e9154429 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -34,6 +34,7 @@ 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 @@ -52,8 +53,8 @@ # 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', 4)) -MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 2)) +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 @@ -63,7 +64,7 @@ MAX_STEPS = int(os.environ.get('MAX_STEPS', 5000)) BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 4)) MINI_BATCH_SIZE = int(os.environ.get('MINI_BATCH_SIZE', 4)) -MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 2)) +MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 1)) GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 1)) SAVE_STEPS = int(os.environ.get('SAVE_STEPS', 100)) @@ -495,7 +496,7 @@ def main(): device_type='GPU'), ] - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS) 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) condenser_mesh = DeviceMesh.from_sizes(world_size=CONDENSER_GPUS, dp_size=CONDENSER_GPUS) @@ -530,6 +531,7 @@ def main(): 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) @@ -590,6 +592,15 @@ def main(): # -- 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 = [] @@ -599,7 +610,7 @@ def prepare_rag_batch(batch): prob = '' for m in msgs: if m.get('role') == 'user': - prob = m.get('content', '') + prob = _extract_text(m.get('content', '')) break problems.append(prob) ud = item.get('user_data', []) From 9cb0662ba3487b7671fa96cb74de583719678dab Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 30 Jun 2026 13:48:22 +0800 Subject: [PATCH 14/26] fix --- cookbook/exp/rl/rag_hint_grpo.py | 346 ++++++++++++++++++++++++++++--- 1 file changed, 314 insertions(+), 32 deletions(-) diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index e9154429..a9e7d840 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -58,15 +58,16 @@ NUM_GPUS = CONDENSER_GPUS + EMB_GPUS + SAMPLER_GPUS + MODEL_GPUS # Training hyperparams -NUM_GENERATIONS = int(os.environ.get('NUM_GENERATIONS', 16)) -MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 65536)) +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', 4)) -MINI_BATCH_SIZE = int(os.environ.get('MINI_BATCH_SIZE', 4)) +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', 1)) 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', 2.0)) # RAG config DB_PATH = os.environ.get('DB_PATH', './output.oldemb/thinking_rag/lance.db') @@ -142,7 +143,8 @@ 'Below are condensed reasoning examples from similar problems. ' 'Analyze them critically — identify which steps/concepts are applicable ' 'and which may not apply. Then solve the actual problem step by step. ' - 'Put your final answer inside \\boxed{{}}.\n\n' + 'Put your final answer inside \\boxed{}. ' + 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.\n\n' ) EXAMPLE_TEMPLATE = ( @@ -154,7 +156,8 @@ SYSTEM_DIRECT = ( 'You are an expert competition mathematician. ' - 'Solve the problem step by step. Put your final answer inside \\boxed{}.' + '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{}.' ) @@ -303,11 +306,23 @@ def extract_boxed(text: str) -> str: 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)=... + _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) @@ -315,15 +330,21 @@ def normalize_answer(ans: str) -> str: s = s.replace(r'\,', '') s = s.replace(r'\;', '') s = s.replace(r'\!', '') - s = s.replace(r'\text', '') - s = s.replace(r'\mathrm', '') + # 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') s = s.strip('$').strip() + # Remove trailing unit braces: {cm}, {kg}, etc. s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|°', 'deg', 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) @@ -350,37 +371,188 @@ def _frac_to_slash(m): 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: - 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 - frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') - def _eval_frac(s): + """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 = _eval_frac(a), _eval_frac(b) + + va, vb = _try_eval(a), _try_eval(b) if va is not None and vb is not None: - return abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + return abs(va - vb) < 1e-6 * max(1, abs(va), abs(vb)) return False + @classmethod + def _normalize_tuple(cls, s: str) -> str: + """Normalize tuple formatting: (2, 5, 609) → 2,5,609.""" + 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 - return cls._try_numeric_equal(norm_p, norm_r) + + # --- 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 + + # --- 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 + + return False def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: rewards = [] @@ -420,13 +592,47 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: 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) - total = [a + f for a, f in zip(acc, fmt)] + gib = gib_fn(trajectories) + total = [a + f + g for a, f, g in zip(acc, fmt, gib)] return total, fmt, acc @@ -496,7 +702,7 @@ def main(): device_type='GPU'), ] - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS) + 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, tp_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) @@ -509,23 +715,23 @@ def main(): 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('GRPOLoss', epsilon=0.2) - model.set_processor(InputProcessor, padding_free=True) + model.set_loss('GRPOLoss', epsilon=0.2, beta=0.04) + model.set_processor(InputProcessor) model.set_template('Qwen3_5Template', model_id=MODEL_ID, - enable_thinking=True, max_length=65536) + enable_thinking=True, max_length=32768) # -- Rollout sampler -- sampler = vLLMSampler( model_id=MODEL_ID, engine_args={ 'gpu_memory_utilization': 0.8, - 'max_model_len': 65536, + '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=65536) + enable_thinking=True, max_length=32768) # -- Embedding model -- emb_model = TransformersModel( @@ -784,6 +990,14 @@ def _fallback(ci): # ---- 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({ @@ -792,7 +1006,6 @@ def _fallback(ci): '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') - rag_log_f.flush() metrics.accumulate( completion_lengths=all_completion_lengths, @@ -806,18 +1019,87 @@ def _fallback(ci): # ---- 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 + # 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 all-same reward problem groups (no gradient 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 + 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 ---- - total_completions = len(all_input_data) + total_completions = len(filtered_inputs) + if total_completions == 0: + logger.info(f'[Step {optim_step}] all groups filtered (uniform rewards), skip training') + continue + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) - mb_inputs = all_input_data[mb_start:mb_end] - mb_old_logps = all_old_logps[mb_start:mb_end] - mb_advantages = advantages[mb_start:mb_end] + 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] model.forward_backward( inputs=mb_inputs, old_logps=mb_old_logps, + ref_logps=mb_old_logps, advantages=mb_advantages, micro_batch_size=MICRO_BATCH_SIZE, ) From 5e8d6f0d2012774433c30c27a745ff7fd429ccac Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 30 Jun 2026 13:57:24 +0800 Subject: [PATCH 15/26] fix --- cookbook/exp/rl/grpo.py | 684 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 684 insertions(+) create mode 100644 cookbook/exp/rl/grpo.py diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/rl/grpo.py new file mode 100644 index 00000000..892ee155 --- /dev/null +++ b/cookbook/exp/rl/grpo.py @@ -0,0 +1,684 @@ +"""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', 1)) +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', 2.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') + 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 _normalize_tuple(cls, s: str) -> str: + 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 '=' 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 + + 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', gpus_per_worker=SAMPLER_GPUS), + 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, tp_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('GRPOLoss', epsilon=0.2, 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, 'a', 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 all-same reward problem groups (no gradient 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 + 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 + total_completions = len(filtered_inputs) + if total_completions == 0: + logger.info(f'[Step {optim_step}] all groups filtered (uniform rewards), skip training') + continue + + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_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] + + model.forward_backward( + inputs=mb_inputs, + old_logps=mb_old_logps, + ref_logps=mb_old_logps, + advantages=mb_advantages, + micro_batch_size=MICRO_BATCH_SIZE, + ) + 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}') + + 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() From a90e7c83235955cc7b6f7e7cb38086c5f208d774 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 30 Jun 2026 15:52:32 +0800 Subject: [PATCH 16/26] fix --- cookbook/exp/rl/grpo.py | 8 ++++---- cookbook/exp/rl/rag_hint_grpo.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/rl/grpo.py index 892ee155..02c6def0 100644 --- a/cookbook/exp/rl/grpo.py +++ b/cookbook/exp/rl/grpo.py @@ -43,8 +43,8 @@ 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)) +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 2)) +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 6)) NUM_GPUS = SAMPLER_GPUS + MODEL_GPUS # Training hyperparams @@ -445,7 +445,7 @@ def main(): device_type='GPU'), ] - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=2) + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=MODEL_GPUS) sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, tp_size=SAMPLER_GPUS) twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, @@ -456,7 +456,7 @@ def main(): 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('GRPOLoss', epsilon=0.2, beta=0.04) + model.set_loss('GSPOLoss', epsilon=0.2, epsilon_high=0.28, beta=0.0) model.set_processor(InputProcessor) model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=True, max_length=32768) diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index a9e7d840..5c8beb17 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -715,7 +715,7 @@ def main(): 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('GRPOLoss', epsilon=0.2, beta=0.04) + model.set_loss('GSPOLoss', epsilon=0.2, epsilon_high=0.28, beta=0.0) model.set_processor(InputProcessor) model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=True, max_length=32768) From 79b76babf7622365813b597a7f7c68a8f9cf708a Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 30 Jun 2026 16:27:22 +0800 Subject: [PATCH 17/26] fix --- cookbook/exp/rl/grpo.py | 36 +++++++++++++++++++++----------- cookbook/exp/rl/rag_hint_grpo.py | 30 ++++++++++++++++++-------- 2 files changed, 45 insertions(+), 21 deletions(-) diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/rl/grpo.py index 02c6def0..b9ea887d 100644 --- a/cookbook/exp/rl/grpo.py +++ b/cookbook/exp/rl/grpo.py @@ -43,8 +43,8 @@ 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', 2)) -MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 6)) +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 @@ -445,7 +445,7 @@ def main(): device_type='GPU'), ] - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=MODEL_GPUS) + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=4) sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, tp_size=SAMPLER_GPUS) twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, @@ -643,14 +643,19 @@ def _content_to_str(content): filtered_old_logps.extend(all_old_logps[g_start:g_end]) filtered_advantages.extend(grp_adv) - # Mini-batch training + # 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 - for mb_start in range(0, total_completions, MINI_BATCH_SIZE): - mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + 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] @@ -660,16 +665,23 @@ def _content_to_str(content): old_logps=mb_old_logps, ref_logps=mb_old_logps, advantages=mb_advantages, - micro_batch_size=MICRO_BATCH_SIZE, ) + accum_count += 1 + + if accum_count % grad_accum_steps == 0: + 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: 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}') - log_dict = metrics.calculate() log_dict.update(model.calculate_metric(is_training=True)) metrics.reset() diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index 5c8beb17..a26d7183 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -1084,14 +1084,19 @@ def _content_to_str(content): filtered_old_logps.extend(all_old_logps[g_start:g_end]) filtered_advantages.extend(grp_adv) - # ---- Mini-batch training ---- + # ---- 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 - for mb_start in range(0, total_completions, MINI_BATCH_SIZE): - mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + 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] @@ -1101,16 +1106,23 @@ def _content_to_str(content): old_logps=mb_old_logps, ref_logps=mb_old_logps, advantages=mb_advantages, - micro_batch_size=MICRO_BATCH_SIZE, ) + accum_count += 1 + + if accum_count % grad_accum_steps == 0: + 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: 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}') - log_dict = metrics.calculate() log_dict.update(model.calculate_metric(is_training=True)) metrics.reset() From 6f9787d423fb5970a92de54d7c30e3ce2c77bb3b Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 30 Jun 2026 16:55:39 +0800 Subject: [PATCH 18/26] fix --- cookbook/exp/rl/grpo.py | 4 ++-- cookbook/exp/rl/rag_hint_grpo.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/rl/grpo.py index b9ea887d..a2299d52 100644 --- a/cookbook/exp/rl/grpo.py +++ b/cookbook/exp/rl/grpo.py @@ -54,7 +54,7 @@ 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', 1)) +MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 2)) 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', 2.0)) @@ -445,7 +445,7 @@ def main(): device_type='GPU'), ] - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=4) + 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, tp_size=SAMPLER_GPUS) twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index a26d7183..e17fea92 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -64,7 +64,7 @@ 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', 1)) +MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 2)) 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', 2.0)) From 637127b8239fcf433a408af3971e9a71a16dd05e Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 30 Jun 2026 19:14:32 +0800 Subject: [PATCH 19/26] fix --- cookbook/exp/rl/grpo.py | 7 ++++++- cookbook/exp/rl/rag_hint_grpo.py | 7 ++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/rl/grpo.py index a2299d52..be1bb46d 100644 --- a/cookbook/exp/rl/grpo.py +++ b/cookbook/exp/rl/grpo.py @@ -631,7 +631,9 @@ def _content_to_str(content): diag_f.flush() - # Filter out all-same reward problem groups (no gradient signal) + # 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 @@ -639,6 +641,9 @@ def _content_to_str(content): 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.1 or grp_acc_rate > 0.9: + 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) diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index e17fea92..4707df43 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -1072,7 +1072,9 @@ def _content_to_str(content): rag_log_f.flush() - # ---- Filter out all-same reward problem groups (no gradient signal) ---- + # ---- 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 @@ -1080,6 +1082,9 @@ def _content_to_str(content): 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.1 or grp_acc_rate > 0.9: + 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) From 71098cc3a668ed0138006ce122fcc8481f9857f9 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 30 Jun 2026 19:49:12 +0800 Subject: [PATCH 20/26] fix --- cookbook/exp/rl/grpo.py | 4 ++-- cookbook/exp/rl/rag_hint_grpo.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/rl/grpo.py index be1bb46d..3439ca85 100644 --- a/cookbook/exp/rl/grpo.py +++ b/cookbook/exp/rl/grpo.py @@ -440,13 +440,13 @@ def main(): device_groups = [ DeviceGroup(name='sampler', ranks=list(range(sampler_start, model_start)), - device_type='GPU', gpus_per_worker=SAMPLER_GPUS), + 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, tp_size=SAMPLER_GPUS) + 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) diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index 4707df43..c6389209 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -697,13 +697,13 @@ def main(): 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', gpus_per_worker=SAMPLER_GPUS), + 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, tp_size=SAMPLER_GPUS) + 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) From 71f6868d178a15f0e478fdcb86be2a717e230d18 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 30 Jun 2026 23:36:26 +0800 Subject: [PATCH 21/26] fix --- cookbook/exp/rl/grpo.py | 57 +++++++++++++- cookbook/exp/rl/rag_hint_grpo.py | 125 +++++++++++++++++++++++++++++-- 2 files changed, 174 insertions(+), 8 deletions(-) diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/rl/grpo.py index 3439ca85..af676601 100644 --- a/cookbook/exp/rl/grpo.py +++ b/cookbook/exp/rl/grpo.py @@ -101,9 +101,10 @@ def extract_boxed(text: str) -> str: _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_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) + _VAR_PREFIX_RE = re.compile( + r'^(?:[a-zA-Z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)', re.DOTALL) _EQ_RHS_RE = re.compile(r'^.+=\s*(.+)$') @staticmethod @@ -123,6 +124,8 @@ def normalize_answer(ans: str) -> str: 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) @@ -215,9 +218,36 @@ def _try_eval(s: str): 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: - return re.sub(r'[\s()\[\]]', '', s) + # 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: @@ -277,6 +307,11 @@ def answers_match(cls, predicted: str, reference: str) -> bool: 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('=') @@ -305,6 +340,22 @@ def _sort_factors(s): 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]: diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index c6389209..a858c5ce 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -73,7 +73,7 @@ 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.65)) +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') @@ -218,6 +218,48 @@ def _wrap_anchor(text: str) -> List[Dict[str, str]]: ] +_DECONTAM_JUDGE_PROMPT = ( + 'Are these two math problems essentially the SAME problem ' + '(same mathematical question, possibly different notation/language)?\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: @@ -310,10 +352,11 @@ def extract_boxed(text: str) -> str: _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_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)=... - _VAR_PREFIX_RE = re.compile(r'^[a-zA-Z](?:\([^)]*\))?\s*=\s*(.+)', re.DOTALL) + # 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*(.+)$') @@ -336,6 +379,8 @@ def normalize_answer(ans: str) -> str: 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) @@ -440,10 +485,37 @@ def _try_eval(s: str): 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.""" - return re.sub(r'[\s()\[\]]', '', s) + # 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: @@ -519,6 +591,11 @@ def answers_match(cls, predicted: str, reference: str) -> bool: 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: @@ -552,6 +629,22 @@ def _sort_factors(s): 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]: @@ -831,6 +924,28 @@ def prepare_rag_batch(batch): query_vecs = get_embeddings(emb_model, emb_template, problems, EMB_GPUS) retrieved = retrieve_topk(tbl, query_vecs, problems, SIM_THRESHOLD) + # 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 = [] From 0f91fc24461e1208f702c6a0007ccae8c69ccb5d Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 1 Jul 2026 10:33:30 +0800 Subject: [PATCH 22/26] fix --- cookbook/exp/rl/grpo.py | 2 +- cookbook/exp/rl/rag_hint_grpo.py | 11 ++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/rl/grpo.py index af676601..660d199d 100644 --- a/cookbook/exp/rl/grpo.py +++ b/cookbook/exp/rl/grpo.py @@ -507,7 +507,7 @@ def main(): 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.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) diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index a858c5ce..3a8f0d50 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -219,8 +219,11 @@ def _wrap_anchor(text: str) -> List[Dict[str, str]]: _DECONTAM_JUDGE_PROMPT = ( - 'Are these two math problems essentially the SAME problem ' - '(same mathematical question, possibly different notation/language)?\n' + 'We are building a RAG-augmented math training system. Problem A is the test ' + 'question; Problem B was retrieved from a knowledge base. If B is essentially ' + 'the same problem as A (same core math, just different wording/notation/format/' + 'negation/numeric values), showing B\'s solution would leak the answer to A.\n' + 'Would using B\'s solution constitute answer leakage for A?\n' 'Problem A: {prob_a}\n' 'Problem B: {prob_b}\n' 'Answer only YES or NO.' @@ -808,7 +811,7 @@ def main(): 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.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) @@ -923,6 +926,7 @@ def prepare_rag_batch(batch): # 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: @@ -1033,6 +1037,7 @@ def _fallback(ci): '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': use_rag, From 22b7d90804cf901cffaf91a532514ff278d6dac8 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 1 Jul 2026 12:02:07 +0800 Subject: [PATCH 23/26] fix --- cookbook/exp/rl/rag_hint_grpo.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index 3a8f0d50..2374b58b 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -220,10 +220,13 @@ def _wrap_anchor(text: str) -> List[Dict[str, str]]: _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. If B is essentially ' - 'the same problem as A (same core math, just different wording/notation/format/' - 'negation/numeric values), showing B\'s solution would leak the answer to A.\n' - 'Would using B\'s solution constitute answer leakage for A?\n' + '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.' From afbb4508a82014fff36165b428d100c99c923b36 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 1 Jul 2026 15:22:27 +0800 Subject: [PATCH 24/26] wip --- cookbook/exp/rl/rag_hint_grpo.py | 148 +++++++++++++++++++++---------- src/twinkle/loss/grpo.py | 26 ++++-- 2 files changed, 119 insertions(+), 55 deletions(-) diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index 2374b58b..75097604 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -64,10 +64,11 @@ 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', 2)) +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', 2.0)) +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') @@ -100,13 +101,7 @@ OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './outputs/rag_hint_grpo') # Forced analysis prefix appended at the start of assistant response -ANALYSIS_PREFIX = ( - "Let me think step by step. First, I will analyze the example above: " - "identify which steps and concepts are CORRECT and APPLICABLE to this problem, " - "and explicitly discard any reasoning that is WRONG, IRRELEVANT, or based on " - "assumptions that do not hold here. Then I will solve the problem using only " - "the validated useful parts:\n\n" -) +ANALYSIS_PREFIX = '' # ============================================================================ # Condenser prompt (strategy-level extraction) @@ -141,9 +136,15 @@ SYSTEM_WITH_RAG_HEADER = ( 'You are an expert competition mathematician. ' 'Below are condensed reasoning examples from similar problems. ' - 'Analyze them critically — identify which steps/concepts are applicable ' - 'and which may not apply. Then solve the actual problem step by step. ' - 'Put your final answer inside \\boxed{}. ' + 'Before solving the problem, you MUST first output a ... block ' + 'that critically analyzes the provided examples:\n' + '- 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' + '- Conclude with a one-line verdict: "Useful: ..." and "Discard: ..."\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' ) @@ -675,19 +676,36 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: class FormatReward(Reward): - """Reward for having \\boxed{} in the output.""" + """Reward for having \\boxed{} and ... analysis in the output.""" + + _HINT_OPEN_RE = re.compile(r'') + _HINT_CLOSE_RE = re.compile(r'') 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 = msg.get('content', '') 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) + # Only check hint tags for RAG prompts (system contains examples) + is_rag = 'condensed reasoning examples from similar problems' in sys_content + if is_rag: + has_hint_open = bool(self._HINT_OPEN_RE.search(completion)) + has_hint_close = bool(self._HINT_CLOSE_RE.search(completion)) + has_hint = has_hint_open and has_hint_close + # 0.3 for boxed + 0.2 for hint analysis = 0.5 max + reward = (0.3 if has_boxed else 0.0) + (0.2 if has_hint else 0.0) + else: + reward = 0.5 if has_boxed else 0.0 + rewards.append(reward) return rewards @@ -1021,35 +1039,31 @@ def _fallback(ci): idx=eidx, example_query=ex['query'], example_thinking=ex['thinking'])) - sys_content = ''.join(parts) + 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, + 'top_retrieved_query': rets[0]['query'][:200] if rets else '', + 'condensed_len': len(examples[0].get('thinking', '')) if examples else 0, + }) else: - sys_content = SYSTEM_DIRECT - - prompt_feature = { - 'messages': [ - {'role': 'system', 'content': sys_content}, - {'role': 'user', 'content': prob}, - ], - 'user_data': [('ground_truth', ground_truths[i])], - 'assistant_prefix': ANALYSIS_PREFIX if use_rag else '', - } - rag_prompts.append(prompt_feature) - - # Diagnostic record - debug_rec = { - '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': use_rag, - } - if rets: - debug_rec['top_retrieved_query'] = rets[0]['query'][:200] - if examples: - debug_rec['condensed_len'] = len(examples[0].get('thinking', '')) - rag_debug_records.append(debug_rec) + # No hint found — skip this query entirely + continue return rag_prompts, rag_debug_records @@ -1199,14 +1213,15 @@ def _content_to_str(content): # 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): + 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.1 or grp_acc_rate > 0.9: + 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]) @@ -1229,7 +1244,7 @@ def _content_to_str(content): mb_old_logps = filtered_old_logps[mb_start:mb_end] mb_advantages = filtered_advantages[mb_start:mb_end] - model.forward_backward( + outputs = model.forward_backward( inputs=mb_inputs, old_logps=mb_old_logps, ref_logps=mb_old_logps, @@ -1238,8 +1253,26 @@ def _content_to_str(content): accum_count += 1 if accum_count % grad_accum_steps == 0: - model.clip_grad_and_step() - optim_step += 1 + # 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 @@ -1248,8 +1281,25 @@ def _content_to_str(content): # Flush remaining accumulated gradients (incomplete window at tail) if accum_count % grad_accum_steps != 0: - model.clip_grad_and_step() - optim_step += 1 + 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)) diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 4bb71216..7f37c26e 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -76,8 +76,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( @@ -87,9 +87,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 @@ -103,7 +110,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, @@ -327,7 +341,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 From 10aecda8dd203a46583dda8cd8ef307b5c210a35 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Fri, 3 Jul 2026 08:07:48 +0800 Subject: [PATCH 25/26] wip --- cookbook/exp/embedding/eval_gpqa_rag.py | 630 ++++++++++++++++++++---- cookbook/exp/rl/grpo.py | 53 +- cookbook/exp/rl/rag_hint_grpo.py | 211 +++++++- 3 files changed, 771 insertions(+), 123 deletions(-) diff --git a/cookbook/exp/embedding/eval_gpqa_rag.py b/cookbook/exp/embedding/eval_gpqa_rag.py index 59957d2d..8ef9e63a 100644 --- a/cookbook/exp/embedding/eval_gpqa_rag.py +++ b/cookbook/exp/embedding/eval_gpqa_rag.py @@ -5,6 +5,11 @@ - ``rag``: Retrieve top-k thinking traces from LanceDB as 1-shot examples, then solve (8 GPUs: DP=4 embedding + TP=4 vLLM). +Optional ``--hint`` flag (rag mode only): + After retrieval (+ optional condensing), call an API model to pre-analyze + which methods from the traces are applicable, then inject the analysis as + a system-prompt "preanalysis" instead of raw traces. + Only problems with ``metadata.boxed == True`` are used (auto-gradable via ``\\boxed{...}`` extraction). A random subset is sampled for efficiency. @@ -15,6 +20,12 @@ # RAG-augmented (8 GPUs) python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag + # RAG + API hint analysis (recommended) + python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag --hint + + # RAG + condense + hint (full pipeline) + python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag --condense --hint + # Smaller sample for quick test python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct --n 100 @@ -63,6 +74,47 @@ CONDENSE_TEMPERATURE = 0.2 CONDENSE_MAX_TOKENS = 8192 +# -- Hint analysis config ------------------------------------------------------ +HINT_ANALYSIS_MAX_TOKENS = int(os.environ.get('HINT_ANALYSIS_MAX_TOKENS', 400)) +HINT_ANALYSIS_TEMPERATURE = 0.3 + +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}' +) + +_PREANALYSIS_BEFORE = ( + 'You are an expert competition mathematician.\n\n' + 'A methodology analysis from a similar problem:\n\n' +) +_PREANALYSIS_AFTER = ( + '\n\n' + 'This is from a related but different problem — ' + 'some techniques may transfer, others may not. ' + 'Solve the given problem step by step and put your final answer inside \\boxed{}.' +) + + +def build_preanalysis_system(hint_analysis: str) -> str: + """Build system prompt with pre-analyzed hint.""" + return _PREANALYSIS_BEFORE + hint_analysis + _PREANALYSIS_AFTER + # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- @@ -142,6 +194,7 @@ def _is_truncated_compression(text: str) -> bool: def _api_throttle(): _api_semaphore.acquire() + wait = 0.0 try: with _api_bucket_lock: now = time.monotonic() @@ -154,9 +207,10 @@ def _api_throttle(): else: wait = (1.0 - _api_tokens[0]) * CONDENSE_API_MIN_INTERVAL _api_tokens[0] = 0.0 - time.sleep(wait) finally: _api_semaphore.release() + if wait > 0: + time.sleep(wait) def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Optional[str]: @@ -177,6 +231,140 @@ def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Opti 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 = 4000 + 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() + return idx, content if content else None + 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], @@ -391,15 +579,43 @@ def _eval_frac(s): return False +# MCQ reference 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 answers_match(predicted: str, reference: str) -> bool: - """Check if two math answers are equivalent.""" + """Check if two math answers are equivalent. + + Supports bidirectional MCQ matching: if reference contains both a letter + and a value (e.g. '\\text{(D) }49'), predicted can match either the letter + or the value. + """ 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 - return _try_numeric_equal(norm_p, norm_r) + if _try_numeric_equal(norm_p, norm_r): + return True + # Bidirectional MCQ matching: reference has letter+value compound format + ref_stripped = reference.strip() + mcq_m = _MCQ_REF_RE.match(ref_stripped) + if mcq_m: + ref_letter = mcq_m.group(1) or mcq_m.group(3) # from either branch + ref_value = (mcq_m.group(2) or mcq_m.group(4) or '').strip() + # predicted is the letter + if norm_p == ref_letter: + return True + # predicted is the numeric value + if ref_value: + norm_rv = normalize_answer(ref_value) + if norm_p == norm_rv or _try_numeric_equal(norm_p, norm_rv): + return True + return False # --------------------------------------------------------------------------- @@ -467,6 +683,16 @@ def build_direct_prompt(problem: str) -> Dict[str, Any]: } +def build_hint_prompt(problem: str, hint_analysis: str) -> Dict[str, Any]: + """Build prompt with pre-analyzed hint in system, problem as user.""" + return { + 'messages': [ + {'role': 'system', 'content': build_preanalysis_system(hint_analysis)}, + {'role': 'user', 'content': problem}, + ] + } + + def build_rag_prompt(problem: str, examples: List[Dict[str, str]]) -> Dict[str, Any]: """Approach B: multi-turn assistant format. @@ -553,11 +779,15 @@ def retrieve_examples(tbl, query_vecs: np.ndarray, top_k: int, decontam_threshold: float = 0.0, ) -> List[List[Dict[str, str]]]: thinking_field = 'thinking_raw' if use_thinking_raw else 'cot_compressed' - # Fetch extra candidates when decontamination is active fetch_limit = top_k + 50 if decontam_threshold > 0 else top_k - all_examples: List[List[Dict[str, str]]] = [] + n_queries = len(query_vecs) + all_examples: List[List[Dict[str, str]]] = [None] * n_queries decontam_skipped = 0 - for qi, vec in enumerate(query_vecs): + _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') @@ -567,6 +797,7 @@ def retrieve_examples(tbl, query_vecs: np.ndarray, top_k: int, ) problem_text = problems[qi] if problems else '' examples = [] + local_skipped = 0 for r in results: if len(examples) >= top_k: break @@ -577,15 +808,21 @@ def retrieve_examples(tbl, query_vecs: np.ndarray, top_k: int, t = r.get(thinking_field, '') if not t: continue - # 13-gram decontamination: skip if retrieved query overlaps with problem if decontam_threshold > 0 and problem_text and q: ng_sim = _ngram_jaccard(problem_text, q) if ng_sim > decontam_threshold: - decontam_skipped += 1 + local_skipped += 1 continue examples.append({'query': q, 'thinking': t, '_sim': round(sim, 4), '_raw_trace_len': len(t)}) - all_examples.append(examples) + 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})') @@ -600,20 +837,31 @@ def main(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument('--mode', choices=['direct', 'rag'], default='direct') - p.add_argument('--n', type=int, default=500, - help='Number of problems to sample (0 = all boxed).') - p.add_argument('--db-path', default='./output/thinking_rag/lance.db') + 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=200, + help='Stop after this many problems are successfully evaluated ' + '(RAG mode: problems that have valid traces after decontam; ' + 'direct mode: ignored, evaluates all filtered problems).') + 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=3) 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) - p.add_argument('--rag-fallback-sim', type=float, default=0.70, - help='Fallback to direct prompt when best retrieval similarity ' - 'is below this threshold (avoids weak-trace loops).') + p.add_argument('--sim-threshold', type=float, default=0.80, + 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', help='Enable condenser re-compression on retrieved traces.') @@ -621,11 +869,25 @@ def main(): 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=True, + help='Enable API hint analysis on retrieved traces (default ON). ' + 'In rag mode: retrieve → condense → API hint analysis → preanalysis system prompt. ' + 'In direct mode: ignored (no traces to analyze).') + p.add_argument('--no-hint', dest='hint', action='store_false', + help='Disable API hint analysis; inject condensed trace directly.') + p.add_argument('--problem-ids-file', default='./output/thinking_rag/aops_rag_problem_ids.json', + 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).') + 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() if args.output is None: - args.output = f'./output/thinking_rag/aops_{args.mode}_results.jsonl' + suffix = f'{args.mode}_hint' if (args.hint and args.mode == 'rag') else args.mode + args.output = f'./output/thinking_rag/aops_{suffix}_results.jsonl' if args.condense and args.use_cot_compressed: logger.warning('--condense requires thinking_raw, ignoring --use-cot-compressed') @@ -635,6 +897,27 @@ def main(): 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 if is_rag: @@ -739,84 +1022,186 @@ def main(): 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') - for batch_start in range(0, len(records), args.batch_size): + # 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) + + 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)) - if is_rag: - 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) - # Strip structural markers from cot_compressed field - if args.use_cot_compressed: - for exs in all_examples: - for ex in exs: - ex['thinking'] = _strip_condenser_markers(ex['thinking']) - - # Condenser re-compression - 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) - - prompts = [] - for r, examples in zip(batch, all_examples): - if not examples: - prompts.append(build_direct_prompt(r['problem'])) - else: - # Drop traces exceeding max_trace_len instead of truncating - filtered = [{'query': ex['query'], 'thinking': ex['thinking']} - for ex in examples - if len(ex['thinking']) <= args.max_trace_len] - # Fallback to direct if best similarity is too low - best_sim = max(ex.get('_sim', 0.0) for ex in examples) - if filtered and best_sim >= args.rag_fallback_sim: - prompts.append(build_rag_prompt(r['problem'], filtered)) - else: - prompts.append(build_direct_prompt(r['problem'])) - else: - prompts = [build_direct_prompt(r['problem']) for r in batch] + return prompts, kept_batch, kept_examples, kept_hints, kept_global_indices, batch_skipped - responses = sampler.sample(prompts, gen_params) + target_reached = False + batch_starts = list(range(0, len(records), args.batch_size)) - 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 - - debug_rec = { - 'idx': batch_start + i, - 'reference_answer': rec['reference_answer'], - 'predicted': predicted, - 'is_correct': is_correct, - 'problem': rec['problem'], - 'model_output': raw_output, - } - if is_rag: + 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, + } debug_rec['num_traces'] = len(all_examples[i]) if all_examples[i]: ex0 = all_examples[i][0] @@ -826,25 +1211,90 @@ def main(): debug_rec['condensed_trace'] = ex0['thinking'] debug_rec['condensed_trace_len'] = len(ex0['thinking']) debug_rec['condense_source'] = ex0.get('_condense_source', '') - debug_records.append(debug_rec) - out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') - out_f.flush() + 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 - acc = correct_count / total_count if total_count else 0 - sys.stderr.write( - f' [{batch_end}/{len(records)}] ' - f'acc={acc:.4f} ({correct_count}/{total_count})\n') + 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, + } + 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'AoPS Math — 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})') 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/rl/grpo.py b/cookbook/exp/rl/grpo.py index 660d199d..7f43d93f 100644 --- a/cookbook/exp/rl/grpo.py +++ b/cookbook/exp/rl/grpo.py @@ -54,10 +54,11 @@ 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', 2)) +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', 2.0)) +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') @@ -550,7 +551,7 @@ def main(): # -- Diagnostics -- os.makedirs(OUTPUT_DIR, exist_ok=True) diag_path = os.path.join(OUTPUT_DIR, 'diagnostics.jsonl') - diag_f = open(diag_path, 'a', encoding='utf-8') + diag_f = open(diag_path, 'w', encoding='utf-8') logger.info(f'[diag] diagnostics → {diag_path}') def _content_to_str(content): @@ -693,7 +694,7 @@ def _content_to_str(content): 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.1 or grp_acc_rate > 0.9: + 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]) @@ -716,7 +717,7 @@ def _content_to_str(content): mb_old_logps = filtered_old_logps[mb_start:mb_end] mb_advantages = filtered_advantages[mb_start:mb_end] - model.forward_backward( + outputs = model.forward_backward( inputs=mb_inputs, old_logps=mb_old_logps, ref_logps=mb_old_logps, @@ -725,8 +726,25 @@ def _content_to_str(content): accum_count += 1 if accum_count % grad_accum_steps == 0: - model.clip_grad_and_step() - optim_step += 1 + 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 @@ -735,8 +753,25 @@ def _content_to_str(content): # Flush remaining accumulated gradients (incomplete window at tail) if accum_count % grad_accum_steps != 0: - model.clip_grad_and_step() - optim_step += 1 + 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)) diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py index 75097604..b35b3706 100644 --- a/cookbook/exp/rl/rag_hint_grpo.py +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -103,6 +103,13 @@ # 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) # ============================================================================ @@ -135,19 +142,71 @@ # ============================================================================ SYSTEM_WITH_RAG_HEADER = ( 'You are an expert competition mathematician. ' - 'Below are condensed reasoning examples from similar problems. ' - 'Before solving the problem, you MUST first output a ... block ' - 'that critically analyzes the provided examples:\n' - '- Identify which steps, formulas, and concepts are CORRECT and APPLICABLE ' - 'to the current problem.\n' + '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' - '- Conclude with a one-line verdict: "Useful: ..." and "Discard: ..."\n\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' @@ -185,6 +244,59 @@ def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Opti _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 # ============================================================================ @@ -678,8 +790,29 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: class FormatReward(Reward): """Reward for having \\boxed{} and ... analysis in the output.""" - _HINT_OPEN_RE = re.compile(r'') - _HINT_CLOSE_RE = re.compile(r'') + _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 = [] @@ -689,20 +822,33 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: sys_content = '' for msg in messages: if msg.get('role') == 'system': - sys_content = msg.get('content', '') + sys_content = self._to_text(msg.get('content', '')) for msg in reversed(messages): if msg.get('role') == 'assistant': - completion = msg.get('content', '') + 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: - has_hint_open = bool(self._HINT_OPEN_RE.search(completion)) - has_hint_close = bool(self._HINT_CLOSE_RE.search(completion)) - has_hint = has_hint_open and has_hint_close - # 0.3 for boxed + 0.2 for hint analysis = 0.5 max - reward = (0.3 if has_boxed else 0.0) + (0.2 if has_hint else 0.0) + # 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) @@ -1023,6 +1169,11 @@ def _fallback(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 = [] @@ -1033,13 +1184,18 @@ def _fallback(ci): use_rag = bool(examples) and best_sim >= RAG_FALLBACK_SIM if use_rag: - 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) + # 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({ @@ -1058,6 +1214,8 @@ def _fallback(ci): '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, }) @@ -1070,7 +1228,7 @@ def _fallback(ci): # 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, 'a', encoding='utf-8') + rag_log_f = open(rag_log_path, 'w', encoding='utf-8') logger.info(f'[rag] diagnostics → {rag_log_path}') batch_iter = iter(dataloader) @@ -1108,6 +1266,10 @@ def _fallback(ci): 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() @@ -1182,7 +1344,8 @@ def _content_to_str(content): 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 + 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 From d64d96dc2c86e84c18dc0e728976faced099d53a Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Sat, 4 Jul 2026 17:48:39 +0800 Subject: [PATCH 26/26] wip --- cookbook/exp/embedding/compare_math_levels.py | 91 ++++ cookbook/exp/embedding/eval_gpqa_rag.py | 406 +++++++++++++----- cookbook/exp/embedding/eval_math_by_level.sh | 59 +++ 3 files changed, 448 insertions(+), 108 deletions(-) create mode 100644 cookbook/exp/embedding/compare_math_levels.py create mode 100755 cookbook/exp/embedding/eval_math_by_level.sh 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/eval_gpqa_rag.py b/cookbook/exp/embedding/eval_gpqa_rag.py index 8ef9e63a..dd71590c 100644 --- a/cookbook/exp/embedding/eval_gpqa_rag.py +++ b/cookbook/exp/embedding/eval_gpqa_rag.py @@ -1,42 +1,41 @@ -"""AoPS math competition evaluation: direct vs RAG-augmented with Qwen3.5-4B. +"""Math evaluation: direct vs RAG-augmented with Qwen3.5-4B. -Two modes: +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``: Retrieve top-k thinking traces from LanceDB as 1-shot - examples, then solve (8 GPUs: DP=4 embedding + TP=4 vLLM). + - ``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 (+ optional condensing), call an API model to pre-analyze - which methods from the traces are applicable, then inject the analysis as - a system-prompt "preanalysis" instead of raw traces. + After retrieval + condensing, call an API model to filter and refine the + traces — keeping only applicable methods — then inject the refined trace. -Only problems with ``metadata.boxed == True`` are used (auto-gradable via -``\\boxed{...}`` extraction). A random subset is sampled for efficiency. +Reference answers are the ``\\boxed{...}`` content of each solution. Launch examples: - # Direct (4 GPUs, default 500 problems) - python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct - - # RAG-augmented (8 GPUs) - python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag + # Default: raw RAG on MATH, stratified 100/level (needs COMPRESS_API_KEY) + COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py - # RAG + API hint analysis (recommended) - python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag --hint - - # RAG + condense + hint (full pipeline) - python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag --condense --hint - - # Smaller sample for quick test - python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct --n 100 + # Paired direct baseline on the same MATH subset + python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct - # RAG with condenser (retrieves thinking_raw, compresses with local 4B / API) - EVAL_CONDENSER_GPUS=2 python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag --condense + # Raw RAG without condenser (inject raw retrieved traces) + python cookbook/exp/embedding/eval_gpqa_rag.py --no-condense - # RAG without condenser (uses thinking_raw by default, truncated to max-trace-len) - python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag + # Add hint filtering back on top of condensing + COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py --hint - # RAG with pre-compressed field (opt-in, not recommended for reader LM) - python cookbook/exp/embedding/eval_gpqa_rag.py --mode rag --use-cot-compressed + # Fall back to the old AoPS dataset + python cookbook/exp/embedding/eval_gpqa_rag.py --dataset aops """ import argparse import json @@ -75,46 +74,31 @@ CONDENSE_MAX_TOKENS = 8192 # -- Hint analysis config ------------------------------------------------------ -HINT_ANALYSIS_MAX_TOKENS = int(os.environ.get('HINT_ANALYSIS_MAX_TOKENS', 400)) -HINT_ANALYSIS_TEMPERATURE = 0.3 +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 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' + '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. 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."' + '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' - '## Condensed Trace (from similar problem)\n{thinking}' -) - -_PREANALYSIS_BEFORE = ( - 'You are an expert competition mathematician.\n\n' - 'A methodology analysis from a similar problem:\n\n' -) -_PREANALYSIS_AFTER = ( - '\n\n' - 'This is from a related but different problem — ' - 'some techniques may transfer, others may not. ' - 'Solve the given problem step by step and put your final answer inside \\boxed{}.' + '## Retrieved Reasoning Traces\n{thinking}' ) - -def build_preanalysis_system(hint_analysis: str) -> str: - """Build system prompt with pre-analyzed hint.""" - return _PREANALYSIS_BEFORE + hint_analysis + _PREANALYSIS_AFTER - # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- @@ -134,6 +118,7 @@ def build_preanalysis_system(hint_analysis: str) -> str: 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') # --------------------------------------------------------------------------- @@ -237,7 +222,7 @@ def _api_hint_analysis_batch( condensed_examples: List[List[Dict[str, str]]], ) -> List[Optional[str]]: """Call API to pre-analyze RAG relevance for each problem.""" - _MAX_HINT_INPUT = 4000 + _MAX_HINT_INPUT = 8000 results: List[Optional[str]] = [None] * len(problems) tasks = [] for i, prob in enumerate(problems): @@ -266,7 +251,10 @@ def _call_one(idx, msgs): 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 + # 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 @@ -373,6 +361,7 @@ def condense_traces( 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. @@ -404,11 +393,23 @@ def condense_traces( 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 = '' @@ -509,8 +510,8 @@ def normalize_answer(ans: str) -> str: if not ans: return '' s = ans.strip() - # MCQ: extract bare letter from \textbf{(D)}, \text{(A)}, (B), etc. - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf)?\{?\(?([A-E])\)?\}?$', s) + # 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(' ', '') @@ -579,19 +580,40 @@ def _eval_frac(s): return False -# MCQ reference pattern: \text{(D) }49, \textbf{(C)}12, (B) 21, etc. +# 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: if reference contains both a letter - and a value (e.g. '\\text{(D) }49'), predicted can match either the letter - or the value. + 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 @@ -601,20 +623,29 @@ def answers_match(predicted: str, reference: str) -> bool: return True if _try_numeric_equal(norm_p, norm_r): return True - # Bidirectional MCQ matching: reference has letter+value compound format - ref_stripped = reference.strip() - mcq_m = _MCQ_REF_RE.match(ref_stripped) - if mcq_m: - ref_letter = mcq_m.group(1) or mcq_m.group(3) # from either branch - ref_value = (mcq_m.group(2) or mcq_m.group(4) or '').strip() - # predicted is the letter - if norm_p == ref_letter: + + # 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 - # predicted is the numeric value - if ref_value: - norm_rv = normalize_answer(ref_value) - if norm_p == norm_rv or _try_numeric_equal(norm_p, norm_rv): - return True return False @@ -649,6 +680,74 @@ def load_aops(n: int, seed: int = 42) -> List[Dict[str, Any]]: 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 # --------------------------------------------------------------------------- @@ -673,24 +772,54 @@ def load_aops(n: int, seed: int = 42) -> List[Dict[str, Any]]: '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}, + {'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 system, problem as user.""" - return { - 'messages': [ - {'role': 'system', 'content': build_preanalysis_system(hint_analysis)}, - {'role': 'user', 'content': problem}, - ] - } + """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, @@ -710,7 +839,7 @@ def build_rag_prompt(problem: str, 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}) + messages.append({'role': 'user', 'content': RAG_FOLLOWUP + MCQ_INSTRUCTION}) return {'messages': messages} @@ -836,20 +965,31 @@ def _search_one(qi: int): def main(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--mode', choices=['direct', 'rag'], default='direct') + 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=200, + p.add_argument('--target-eval', type=int, default=0, help='Stop after this many problems are successfully evaluated ' - '(RAG mode: problems that have valid traces after decontam; ' - 'direct mode: ignored, evaluates all filtered problems).') + '(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=3) + 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.80, + 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, @@ -863,37 +1003,52 @@ def main(): 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', - help='Enable condenser re-compression on retrieved traces.') + 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=True, - help='Enable API hint analysis on retrieved traces (default ON). ' - 'In rag mode: retrieve → condense → API hint analysis → preanalysis system prompt. ' - 'In direct mode: ignored (no traces to analyze).') + 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 analysis; inject condensed trace directly.') - p.add_argument('--problem-ids-file', default='./output/thinking_rag/aops_rag_problem_ids.json', + 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).') + '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/aops_{suffix}_results.jsonl' + 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 - records = load_aops(n=args.n, seed=args.seed) + 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') @@ -920,6 +1075,18 @@ def main(): 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 = [ @@ -1093,7 +1260,8 @@ def _prepare_rag_batch(batch_start: int): condenser_sampler=condenser_sampler_obj, compress_params=condenser_params, special_tokens=condenser_special_tokens, - max_output_len=args.condense_max_len) + max_output_len=args.condense_max_len, + dp_size=condenser_gpus) hint_analyses = None if args.hint and hint_api_client: @@ -1202,6 +1370,10 @@ def _prepare_rag_batch(batch_start: int): '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] @@ -1265,6 +1437,10 @@ def _prepare_rag_batch(batch_start: int): '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() @@ -1276,13 +1452,27 @@ def _prepare_rag_batch(batch_start: int): overall_acc = correct_count / total_count if total_count else 0 print(f'\n{"=" * 60}') - print(f'AoPS Math — mode={args.mode}, model={GEN_MODEL_ID}') + 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}') 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 "============================================================"